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

【问题标题】：Vectorizing Numpy 3D and 2D array operation向量化 Numpy 3D 和 2D 数组操作
【发布时间】：2019-09-11 02:42:14
【问题描述】：

我正在尝试在 Python 中创建K MxN matrices，存储在(M,N,K) numpy 数组C 中，来自两个矩阵A 和B，形状为(K, M) 和@987654328 @ 分别。第一个矩阵计算为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 循环。也许我可以对代码进行矢量化，但我看不到它是如何完成的。

【问题讨论】：

@ 处理批次，但批次维度 k 必须在前面。瞄准 (k,m,n) 并在之后转置。
@hpaulj 请发表您的评论作为答案。

标签： python numpy multidimensional-array vectorization dot-product

【解决方案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

执行时间差不多

【讨论】：