向量化Numpy 3D和2D陣列操作

Question

我正在嘗試在Python中創建K MxN matrices ， A矩陣從兩個矩陣A和B存儲在(M,N,K) numpy數組C ，形狀分別為(K, M)和(K,N) 。 計算第一個矩陣為C0 = a0.T x b0 ，其中a0是A的第一行， b1是B的第一行，第二個矩陣為C1 = a1.T x b0 ，依此類推。

現在，我正在使用for循環來計算矩陣。

import numpy as np
A = np.random.random((10,800))
B = np.random.random((10,500))
C = np.zeros((800,500,10))
for k in range(10):
    C[:,:,k] = A[k,:][:,None] @ B[k,:][None,:]

由於操作是獨立的，所以我想知道是否存在一些避免for循環的pythonic方法。 也許我可以對代碼進行矢量化處理，但是看不到如何完成。

Answer 1

In [235]: A = np.random.random((10,800)) 
     ...: B = np.random.random((10,500)) 
     ...: C = np.zeros((800,500,10)) 
     ...: for k in range(10): 
     ...:     C[:,:,k] = A[k,:][:,None] @ B[k,:][None,:] 
     ...:                                                                            
In [236]: C.shape                                                                    
Out[236]: (800, 500, 10)

批處理基質產品，然后轉置

In [237]: np.allclose((A[:,:,None]@B[:,None,:]).transpose(1,2,0), C)                 
Out[237]: True

但是由於矩陣乘積軸的大小為1，並且沒有其他求和，所以廣播乘法也一樣好：

In [238]: np.allclose((A[:,:,None]*B[:,None,:]).transpose(1,2,0), C)                 
Out[238]: True

執行時間差不多