如何向量化多个矩阵乘法答案

【问题标题】：How to vectorize multiple matrix multiplication如何向量化多个矩阵乘法
【发布时间】：2021-10-07 03:47:17
【问题描述】：

我有一个二维矩阵 A[1000*90] 和 B[90*90*1000]

我要计算C[1000*90]


For i in range(1000)
  C[i,:]=np.matmul(A[i,:],B[:,:,i]

我知道如果我使用矢量化公式会更快，似乎 einsum 可能是我正在寻找的函数，但我在输入 einsum 的语法时遇到了麻烦。是np.einsum(ij,jki->ik,A,B)吗？

【问题讨论】：

你是对的，einsum 就是你想要的。如果你的尝试是正确的（对我来说看起来不错），我自己还不够专业，无法通过眼睛告诉你。但是您可以通过一些小例子轻松说服自己...
matmul 处理“批量”乘法。您在文档中没有看到吗？您可能需要添加维度和转置，因此 i 是 3 个维度中的第一个。
@Walterwang201112 IMO 你的einsum 看起来不错。你为什么不“试试看”？有问题吗？
您的 einsum 是正确的，但 @hpaulj 方式更有效，因为它直接调度到链接的 BLAS 库（如果有）。

标签： python numpy vectorization numpy-einsum

【解决方案1】：

您的 einsum 是正确的。但正如 hpaulj 所指出的，还有一种更好的方法。

使用 Matmul：

import numpy as np
A =np.random.rand(1000,90)
B = np.random.rand(90,90,1000)
C = A[:,np.newaxis,:]@B.transpose(2,0,1) ## Matrix multiplication
C = C = C.reshape(-1,C.shape[2])
np.array_equal(C,np.einsum('ij,jki->ik',A,B)) # check if both give same result

【讨论】：

谢谢，我不得不说 Matmul 真的很难缠住我的脑袋。我不知道你为什么要向 A 添加一个新轴，转置 B，并按照你的方式重塑 C，尤其是使用 3d 矩阵 matmul。在性能方面，两种方法大致相同，matmul 对较小的矩阵稍快，而 einsum 对于较大的矩阵更快。
Einsum 更快，因为您没有将 BLAS 库链接到 Numpy。您可以在以下链接中查看它stackoverflow.com/questions/21671040/… 在处理更大的矩阵时它肯定会给您带来很好的提升。您需要对每个数组执行 1000 个矩阵的乘法运算。为此，您必须匹配 A 和 B 的尺寸，以便 matmul 为您执行 for 循环。 A 的维数为 2，B 为 3。为了使其相同，我添加了一个新维数，以便在 A 中有 1000 个（1,90）矩阵，在 B 中有 1000 个（90,90）矩阵
使用 matmul 时，矩阵乘法是在最后两个维度上执行的，我必须转置 B 以便 1000 成为我的第一个维度。现在， A[:,newaxis,:] 有 (1000,1,90) 并且 B.transpose(2,0,1) 有 (1000,90,90) 维度。现在您可以直接执行 matmul 以便它自动为您执行超过 1000 次的 for 循环。在 matmul 之后，维度为 (1000,1,90)。我对其进行了重塑以使其成为（1000,90）