如何有效地计算numpy中两个矩阵的外积？

Question

假设我有A（KxMxN）和B（KxLxN）矩阵，其中L，M，N很小而K是一个大数。我想计算沿第一维使用最后2维的外积来得到矩阵C（KxMxL）。

我可以通过为＆＃34; K＆＃34;中的每个索引 k 运行for循环来做到这一点。并使用numpy的matmul函数进行2D矩阵

out = [np.matmul(x,y.T) for x, y in zip(A, B)]
out=np.asarray(out)

我想知道如果没有for循环/理解我是否可以这样做，因为K是一个非常大的数字。

Answer 1

由于A的形状为(K, M, N)而B的形状为(K, L, N)，并且您希望找到形状为(K, M, L)的产品总和，您可以使用np.einsum：

C = np.einsum('kmn,kln->kml', A, B)

Answer 2

matmul适用于B中的转置，因此其倒数第二个与A的最后一个匹配。

In [1019]: A=np.random.rand(K,M,N)
In [1021]: B=np.random.rand(K,L,N)

In [1023]: C=np.einsum('kmn,kln->kml',A,B)
In [1024]: C.shape
Out[1024]: (2, 4, 3)

In [1026]: D=A@B.transpose(0,2,1)
In [1027]: D.shape
Out[1027]: (2, 4, 3)

In [1028]: np.allclose(C,D)
Out[1028]: True

对于这个小例子，timeit是相同的。

[np.dot(x,y.T) for x, y in zip(A, B)]做同样的事情;将y的第二个最后一个调整与x的调整相匹配，并在A和B的第一个调暗点上进行迭代。