如何向量化多个矩阵乘法

Question

我有一个二维矩阵 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)吗？

Answer 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

如何向量化多个矩阵乘法

How to vectorize multiple matrix multiplication

python

numpy

vectorization

numpy-einsum