Numpy多维(3d)矩阵乘法

Question

我得到了两个三维矩阵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))执行矩阵乘法。