Numpy ndarray 乘法答案

【问题标题】：Numpy ndarray multiplicationNumpy ndarray 乘法
【发布时间】：2017-02-01 17:42:40
【问题描述】：

我有两个 3D numpy ndarray

A=np.array([[[1, 1],
             [1, 1],
             [1, 1]],

            [[2, 2],
             [2, 2],
             [2, 2]]])

B=np.array([[[ 2,  0],
             [ 0,  2]],

            [[ 2, -2],
             [-2,  2]]])

我想用元素 AB_ijk=∑_m (A_ijm*B_imk) 创建 AB 数组>) 其中求和仅在 m-index 上（重复）而不是在 i 上（依次重复）。

换句话说，我可以通过这个 for 循环获得 di AB ndarray

for i in range(2):
    AB[i,:,:]=np.dot(A[i,:,:],B[i,:,:])

AB等于

array([[[ 2.,  2.],
    [ 2.,  2.],
    [ 2.,  2.]],

   [[ 0.,  0.],
    [ 0.,  0.],
    [ 0.,  0.]]])

有没有办法避免 for 循环？ tensordot或einsum如何获取AB数组？

谢谢你的回答，我真的很感激。

【问题讨论】：

标签： python arrays numpy

【解决方案1】：

在足够新的 NumPy (1.10+) 上，您可以这样做

AB = np.matmul(A, B)

或者（如果你也有 Python 3.5+）：

AB = A @ B

如果你没有 NumPy 1.10+，你可以这样做

AB = np.einsum('ijm,imk->ijk', A, B)

对于较大的 J/M/K 尺寸，特别是如果您有良好的 BLAS，可能还值得考虑使用 dot 的显式 for 循环。与更多解释 Python 所损失的开销相比，BLAS 矩阵乘法可能会节省更多时间。我认为np.matmul 和@ 应该利用dot 所做的相同的事情，但我不认为np.einsum 这样做。

【讨论】：

已确认：einsum 尚未使用 BLAS。

【解决方案2】：

ABijk=∑m (Aijm*Bimk) 转换为

AB = np.einsum('ijm,imk->ijk', A, B)

我认为matmul 运营商也会处理这个问题

AB = A @ B

因为它在最后 2 个维度上使用普通的dot，所以将其余的作为免费行李随身携带。

测试它们，让我知道它们是否有效。

【讨论】：