向量化大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。下面是使用np.einsum做第一张是没有任何额外的内存开销一个vectorized方法 -

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

Inst你可以使用tensordot的dot。您当前的方法等效于：

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

注意transpose末把轴以正确的顺序。

与dot类似，tensordot函数可以调用快速BLAS/LAPACK库（如果已安装它们），所以对于大型数组应该可以很好地执行。