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)）。