我有两个形状为(N,2,2)的3d数组A和B,我想根据N轴将每个2x2矩阵上的矩阵乘积逐元素相乘。通过循环实现,它看起来像
C[i] = dot(A[i], B[i])
有没有办法不用循环就可以做到这一点?我研究过Tensordot,但无法使其正常工作。我想我可能想要
tensordot(a, b, axes=([1,2], [2,1]))
之类的东西,但这给了我NxN矩阵。 最佳答案
似乎您正在沿第一个轴对每个切片进行矩阵乘法。同样,您可以像这样使用 np.einsum
-
np.einsum('ijk,ikl->ijl',A,B)
我们也可以使用
np.matmul
-np.matmul(A,B)
在Python 3.x上,此
matmul
操作简化为 @
operator-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 loop
100 loops, best of 3: 3.04 ms per loop
100 loops, best of 3: 10.9 ms per loop
关于python - 3D数组的块状元素乘积,我们在Stack Overflow上找到一个类似的问题:https://stackoverflow.com/questions/31957364/