Numpy 3D矩阵乘法

Question

我有2个矩阵A（形状10x10x36）和B（形状10x27x36）。我想将最后两个轴相乘并沿轴0求和，以使结果C的形状为10x27。目前，这是我的操作方法

C = []
for i in range(A.shape[0]):
    C.append(np.matmul(A[i], B[i].T))
C = np.sum(np.array(C), axis=0)

我想以向量化的方式实现这一目标，但似乎找不到答案。我已经签出了np.einsum，但不确定如何应用它来获得结果。任何帮助将不胜感激。谢谢！

Answer 1

使用np.einsum获得相同的结果：

r1 = np.einsum('ijk,ilk->jl', A, B)

但是在我的机器中，for循环的执行速度快了将近2倍：

def f(A,B):
    C = []
    for i in range(A.shape[0]):
        C.append(np.matmul(A[i], B[i].T))
    return np.sum(np.array(C), axis=0)

%timeit np.einsum('ijk,ilk->jl',A,B)
102 µs ± 3.79 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%timeit f(A,B)
57.6 µs ± 1.7 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

Answer 2

matmul支持堆叠。您可以简单地做：

 ([email protected](0,2,1)).sum(0)

检查（C是使用OP的循环生成的：]

np.allclose(([email protected](0,2,1)).sum(0),C)
# True
timeit(lambda:([email protected](0,2,1)).sum(0),number=1000)
# 0.03199950899579562
# twice as fast as original loop

Answer 3

您还可以使用列表理解来尝试以下操作。它比您当前使用的要简洁。

C=np.array([A[i] @ B.T[:,:,i] for i in range(10)]).sum(0)