Python 中无循环的三维矩阵乘法

Question

我想在Python（使用numpy库）中做以下操作。

Matrix A is M x N x R
Matrix B is N x 1 x R

我想进行矩阵相乘，AB = C，其中C是一个M x 1 x R的矩阵。简单来说，A中的每个M x N层（总共有R层）都要和B中的每个N x 1的向量独立地进行矩阵相乘。我相信这可以用一行代码实现。我尝试使用tensordot()，但得到的结果似乎和我预期的不一样。

我在Igor Pro编程已经快10年了，现在我想把很多代码转换成Python。

Answer 1

另一种方法（对于像我这样不太熟悉爱因斯坦记法的人来说更简单）是使用 np.matmul()。关键是要确保最后两个维度的大小是匹配的，也就是((M, N) x (N, 1))。为了做到这一点，可以用 np.transpose()。举个例子：

M, N, R = 4, 3, 10
A = np.ones((M, N, R))
B = np.ones((N, 1, R))

# have the matching dimensions at the very end
C = np.matmul(np.transpose(A, (2, 0, 1)), np.transpose(B, (2, 0, 1))) 
C = np.transpose(C, (1, 2, 0))

print(A.shape)
# out: #(4, 3, 10)
print(B.shape)
# out: #(3, 1, 10)
print(C.shape)
# out: #(4, 1, 10)

Answer 2

numpy.tensordot() 是正确的方法：

a = numpy.arange(24).reshape(2, 3, 4)
b = numpy.arange(12).reshape(3, 1, 4)
c = numpy.tensordot(a, b, axes=[1, 0]).diagonal(axis1=1, axis2=3)

编辑：最初的版本有问题，这个版本计算的内容比应该的多，而且大部分结果都被丢掉了。也许用 Python 在最后一个轴上循环会是更好的方法。

另一个编辑：我得出的结论是 numpy.tensordot() 在这里 并不是 最好的解决方案。

c = (a[:,:,None] * b).sum(axis=1)

这样做会更高效（虽然理解起来会更难）。

Answer 3

抱歉打扰了，不过这个答案可以大大改进，使用非常有用的 np.einsum。

import numpy as np

D,M,N,R = 1,2,3,4
A = np.random.rand(M,N,R)
B = np.random.rand(N,D,R)

print np.einsum('mnr,ndr->mdr', A, B).shape

要知道，它有几个优点：首先，它运行得很快。np.einsum 通常优化得很好，而且它聪明到可以避免创建一个 MxNxR 的临时数组，而是直接在 N 上进行计算。

但也许更重要的是，它非常易读。毫无疑问，这段代码是正确的；而且你可以轻松地让它变得更复杂。

如果你愿意，可以直接从 B 和 einsum 语句中去掉虚拟的 'D' 轴。