如何有效地计算出两个矩阵级数的外积？

Question

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

我可以通过对"K“中的每个索引k运行一个for循环来做到这一点，并对2D矩阵使用numpy的matmul函数

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

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

Answer 1

由于A有shape (K, M, N)，B有shape (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的第一个dim上迭代。