最快的方式

Question

我有一个大小为（M，N，N）的三维张量A.我还有一个长度为M的权重向量p。我想计算

维度N可能很大，所以我想以有效的方式实现它。我使用以下代码：

import numpy as np
temp=np.array([p[m]*A[m] for m in range(M)])
B=sum(temp);

我想知道是否有更快更有效的方法来实现这一点。这很重要的原因是因为当N很大时，上述变量temp存储这些N * N矩阵的M.所以我想我们可以使用“for循环”，但我读到for循环比列表理解慢。

Answer 1

至少从我测试过的张量来看，这个片段的速度提高了3倍：

sum(np.einsum('ij...,i->ij...', A, p))

以下是共享的各种方法的比较：

base      = sum(np.array([p[m]*A[m] for m in range(M)]))
einsum1   = sum(np.einsum('ij...,i->ij...', A, p))
einsum2   = np.einsum('ijk,i->jk', A, p)
einsum3   = np.einsum('ij...,i->j...', A, p)
dot       = p.dot(A.reshape(A.shape[0],-1)).reshape(A.shape[1],-1)
tensordot = np.tensordot(p, A, (0, 0))

对于M = 1000，N = 100（与基础相比）：

einsum1   = 1.11x
einsum2   = 5.83x
einsum3   = 6.26x
dot       = 8.60x
tensordot = 9.83x

Answer 2

In [24]: A=np.arange(3*2*2).reshape(3,2,2)    
In [25]: p=np.array([1,2,3])

In [26]: temp=np.array([p[m]*A[m] for m in range(len(p))])
In [27]: sum(temp)  # python sum
Out[27]: 
array([[32, 38],
       [44, 50]])

In [28]: np.einsum('i...,i->...',A,p)
Out[28]: 
array([[32, 38],
       [44, 50]])

ijk,i->jk也有效。

虽然不一定最快，但temp可以通过简单广播来计算。 np.sum(..., axis=0)相当于Python sum()。

In [30]: (A*p[:,None,None]).sum(axis=0)
Out[30]: 
array([[32, 38],
       [44, 50]])