多维矩阵（/ array）的矩阵乘法 - 如何避免循环？

Question

我正在尝试使用包含多个矩阵的数组来计算矩阵乘法。这可以通过使用np.dot（或Py3.5 +中的新@运算符）的两个矩阵轻松实现，但我正在努力扩展它以有效地评估我的多维数组。

作为一个例子，假设我的矩阵A有形状（5,3,3），B有形状（5,3）。现在，我想为每5个案例矩阵多个后面的部分：即做

res[0] = np.dot(A[0], B[0])
res[1] = np.dot(A[1], B[1])
etc

我可以使用循环成功实现这一点 - 例如：

A = np.random.random((5,3,3))
B = np.random.random((5,3))

res = np.zeros([5,3])

for i in range(len(A)):
    res[i] = np.dot(A[i], B[i])

虽然这很慢，因为它使用循环。

我可以采取一种功能/方法来完全矢量化吗？

感谢。

Answer 1

您可以使用np.einsum -

np.einsum('ijk,ik->ij',A,B)

使用np.matmul -

np.matmul(A,B[...,None]).squeeze()
np.matmul(A,B[...,None])[...,0]

Answer 2

另外，使用和和广播：

np.sum(A * B[:, np.newaxis, :], axis=2)

但是，Einsum似乎是最快的选择。 这将是最慢的，因为它实例化元素的产物。

Answer 3

因为我总是很难理解einsum()，所以我想出了这个：

np.diagonal(np.squeeze(np.dot(A, B[:,:,None])), axis2=2).T

它有效，并且它没有循环，但它明显慢于einsum()，因为:,:,None从B从2D扩展到3D（后来减少了{{ 1}}）。我不确定是否有办法让这个更紧凑。