向量化 Numpy 3D 和 2D 数组操作

Question

我正在尝试在 Python 中创建 K MxN matrices，存储在 (M,N,K) numpy 数组 C 中，来自两个矩阵 A 和 B，形状分别为 (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,:]

由于操作是独立的，我想知道是否有一些 pythonic 方法可以避免 for 循环。也许我可以对代码进行矢量化，但我看不出它是如何完成的。

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

执行时间差不多

向量化 Numpy 3D 和 2D 数组操作

Vectorizing Numpy 3D and 2D array operation

python

numpy

multidimensional-array

vectorization

dot-product