向矢量化大NumPy乘法

Question

我感兴趣的是计算一个大型NumPy数组。我有一个大数组A，它包含一串数字。我要计算这些数字的不同组合的和。这些数据的结构如下：

A = np.random.uniform(0,1, (3743, 1388, 3))
Combinations = np.random.randint(0,3, (306,3))
Final_Product = np.array([  np.sum( A*cb, axis=2)  for cb in Combinations])

我的问题是，是否有一种更优雅和内存更有效的方法来计算这一点？当涉及到三维数组时，我发现使用np.dot()是件令人沮丧的事情.

如果有帮助，Final_Product的理想形状应该是(3743,306,1388)。目前Final_Product是形状(306,3743,1388)，所以我可以只是重塑到那里。

Answer 1

np.dot()不会给出所需的输出，除非涉及可能包含reshaping的额外步骤。这里有一种vectorized方法，使用np.einsum来完成它，而不需要任何额外的内存开销-

Final_Product = np.einsum('ijk,lk->lij',A,Combinations)

为了完整起见，下面是前面讨论过的np.dot和reshaping -

M,N,R = A.shape
Final_Product = A.reshape(-1,R).dot(Combinations.T).T.reshape(-1,M,N)

运行时测试和验证输出-

In [138]: # Inputs ( smaller version of those listed in question )
     ...: A = np.random.uniform(0,1, (374, 138, 3))
     ...: Combinations = np.random.randint(0,3, (30,3))
     ...: 

In [139]: %timeit np.array([  np.sum( A*cb, axis=2)  for cb in Combinations])
1 loops, best of 3: 324 ms per loop

In [140]: %timeit np.einsum('ijk,lk->lij',A,Combinations)
10 loops, best of 3: 32 ms per loop

In [141]: M,N,R = A.shape

In [142]: %timeit A.reshape(-1,R).dot(Combinations.T).T.reshape(-1,M,N)
100 loops, best of 3: 15.6 ms per loop

In [143]: Final_Product =np.array([np.sum( A*cb, axis=2)  for cb in Combinations])
     ...: Final_Product2 = np.einsum('ijk,lk->lij',A,Combinations)
     ...: M,N,R = A.shape
     ...: Final_Product3 = A.reshape(-1,R).dot(Combinations.T).T.reshape(-1,M,N)
     ...: 

In [144]: print np.allclose(Final_Product,Final_Product2)
True

In [145]: print np.allclose(Final_Product,Final_Product3)
True

Answer 2

而不是dot，您可以使用tensordot。您的当前方法相当于：

np.tensordot(A, Combinations, [2, 1]).transpose(2, 0, 1)

请注意末尾的transpose以正确的顺序放置轴。

与dot类似，tensordot函数可以调用快速的BLAS/LAPACK库(如果已经安装了它们)，因此对于大型数组应该执行得很好。