沿多维数组指定轴的向量化矩阵-向量乘法答案

【问题标题】：Vectorized Matrix-Vector multiplication along specified axis of multidimensional array沿多维数组指定轴的向量化矩阵-向量乘法
【发布时间】：2021-12-12 22:32:12
【问题描述】：

假设我有一个大小为 (M,N,D,D) 的数组 A 和一个大小为 (M,N,D) 的向量 x。我想生成大小为 (M,N,D) 的数组 b 其中 b[m,n,:] = np.matmul(A[m,n,:,:],x[m,n,:]) .所以基本上，包含在 A[m,n,:] 中的 D-by-D 矩阵和包含在 x[m,n,:] 中的大小为 D 的向量应该在标准矩阵向量乘法之后相乘。一个天真的 for 循环实现是

for m in range(M):
    for n in range(N):
        b[m,n,:] = np.matmul(A[m,n,:],x[m,n,:])

显然这对于大的 N 和 M 来说是低效的（对我来说就是这种情况；N 大约为 100,000，M 大约为 1000）。我的问题是：如何向量化这个循环？我尝试过使用 np.tensordot 但无论我如何尝试构造它，我都会遇到形状不匹配的问题，大概是因为 A 有 4 个轴，而 x 有 3 个（请注意，据我所知， np.tensordot is使用正确的功能，我只是没有正确使用它）

编辑尝试使用 np.einsum，但我没有得到预期的结果（应该是 [5,11]：

【问题讨论】：

np.matmul(A[m,n,:,:],x[m,n,:])对两个数组的最后一个维度进行乘积求和，其中x[m,n,:] 是1d，并由（根据文档）作为x[m,n,:,None]（第二至最后）。 matmul 的关键概念是“A 的最后一个，B 的第二个到最后一个”。主要维度，这里m,n 是“批量”维度。（tensordot 不处理“批次”）。

标签： arrays numpy multidimensional-array scipy vectorization

【解决方案1】：

这就是np.einsum 的亮点：

np.einsum('mnij,mnj->mni', A,B)

检查

np.allclose(np.einsum('mnij,mnj->mni', A,x), b)

返回True

【讨论】：

我试过这个；虽然最终结果具有正确的形状，但它不会产生正确的输出。考虑情况 A[0,0,:] = np.array([[1,2],[3,4]]) 和 x[0,0,:] = np.array([1,2]) .在这种情况下，我希望 b[0,0,:] = np.array([5,11])，因为这是矩阵向量乘积。但是，运行 b = np.einsum('mndd,mnd->mnd',A,x) 表明 b[0,0,:] = np.array([1,8])。
@torola 你应该仔细检查代码。我的系统上有[5, 11]。
我已经编辑了原始帖子以显示我的输出 - 错误可能是由于我如何将值分配给 A 和/或 x？
@torola 我的公式是'mnij,mnj->mni'，而不是mndd,mnd->mnd，这是跟踪公式。
啊，现在说得通了（我认为你使用的 i 和 j 只是占位符，我应该用 d 替换它们）。现在完美运行，非常感谢！

【解决方案2】：

广播也可以达到同样的效果：

x1 = x.reshape(M,N,d,1)
b1 = np.matmul(A,x1)
b2 = b1.reshape(M,N,d)

【讨论】：