似乎您正在沿第一个轴对每个切片进行矩阵乘法。同样，您可以这样使用
np.einsum-
np.einsum('ijk,ikl->ijl',A,B)我们还可以使用
np.matmul-
np.matmul(A,B)
在Python
3.x上，此
matmul操作可简化为
@运算符-
A @ B
标杆管理
方法-
def einsum_based(A,B): return np.einsum('ijk,ikl->ijl',A,B)def matmul_based(A,B): return np.matmul(A,B)def forloop(A,B): N = A.shape[0] C = np.zeros((N,2,2)) for i in range(N): C[i] = np.dot(A[i], B[i]) return C时间-
In [44]: N = 10000 ...: A = np.random.rand(N,2,2) ...: B = np.random.rand(N,2,2)In [45]: %timeit einsum_based(A,B) ...: %timeit matmul_based(A,B) ...: %timeit forloop(A,B)100 loops, best of 3: 3.08 ms per loop100 loops, best of 3: 3.04 ms per loop100 loops, best of 3: 10.9 ms per loop
