numpy的ndarray乘法

我有两个3D numpy的ndarraynumpy的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创建AB阵列_IJK =Σ_米（A _IJM * B _IMK），其中总和只是在m-索引（重复）之上，而不是在i之上（这又是重复的）。

换句话说，我可以得到双AB ndarray这个for循环

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.]]])

有没有办法避免的循环？我如何用tensordot或者einsum获得AB阵列？

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

来源

2017-02-01 musine

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

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

我觉得matmul运营商也将处理这个

AB = A @ B

，因为它需要在最后2个维度正常dot，以及免费行李携带休息。

测试这些，让我知道他们是否工作。

来源

2017-02-01 17:50:58 hpaulj

确认，他们都工作 – musine

在一个足够新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，它也可能是值得考虑的EXPL icit for与dot循环。 BLAS矩阵乘法可能会节省更多的时间，而不是解释更多的Python丢失的开销。我认为np.matmul和@应该利用dot所做的相同的事情，但我不认为np.einsum的确如此。

来源

2017-02-01 17:51:37 user2357112

确认：'einsum'不使用BLAS。 – Daniel

numpy的ndarray乘法

回答

相关问题