Python：具有额外索引的快速矩阵乘法

Question

我有两个数组，A 和 B，维度分别为 (l,m,n) 和 (l,m,n,n)。我想获得一个维度为 (l,m,n) 的数组 C，它是通过将 A 和 B 视为其第四个 (A) 以及第三个和第四个索引 (B) 中的矩阵而获得的。一个简单的方法是：


import numpy as np

#Define dimensions

l = 1024
m = l
n = 6

#Create some random arrays

A = np.random.rand(l,m,n)
B = np.random.rand(l,m,n,n)
C = np.zeros((l,m,n))

#Desired multiplication

for i in range(0,l):

    for j in range(0,m):

        C[i,j,:] = np.matmul(A[i,j,:],B[i,j,:,:])

但是，速度很慢（在我的 MacBook 上大约需要 3 秒）。最快、完全矢量化的方法是什么？

Answer 1

尝试使用einsum.

它有很多用例，查看文档：https://numpy.org/doc/stable/reference/generated/numpy.einsum.html

或者，要了解更多信息，还可以在以下位置找到非常好的解释：https://ajcr.net/Basic-guide-to-einsum/

在你的情况下，似乎

np.einsum('dhi,dhij->dhj',A,B)

应该可以。此外，如果需要，您可以尝试 optimize=True 标志以获得更快的速度。

Python：具有额外索引的快速矩阵乘法

Python: fast matrix multiplication with extra indices

python

python-3.x

performance

matrix

matrix-multiplication