python - 多维矩阵(/数组)的矩阵乘法 - 如何避免循环？

Question

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

例如，假设我有一个形状为 (5,3,3) 的矩阵 A 和一个形状为 (5,3) 的矩阵 B。现在，我想为每 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]