numpy 多维（3d）矩阵乘法

Question

我得到两个 3d 矩阵 A（32x3x3）和 B（32x3x3），我想得到维度为 32x3x3 的矩阵 C。计算可以使用循环完成，如：

a = numpy.random.rand(32, 3, 3)
b = numpy.random.rand(32, 3, 3)
c = numpy.random.rand(32, 3, 3)

for i in range(32):
    c[i] = numpy.dot(a[i], b[i])

我相信必须有一个更有效的单线解决方案来解决这个问题。有没有人可以帮忙，谢谢。

Answer 1

您可以使用以下方法执行此操作np.einsum：

In [142]: old = orig(a,b)

In [143]: new = np.einsum('ijk,ikl->ijl', a, b)

In [144]: np.allclose(old, new)
Out[144]: True

使用的一个优点einsum是您几乎可以从索引中读出它正在执行的操作：单独保留第一个轴 ( i)，并对最后两个轴执行矩阵乘法(jk,kl->jl)。