向量化大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庫（如果已安裝它們），所以對於大型數組應該可以很好地執行。