Numpy 3D 矩陣乘法

Question

我有 2 個矩陣 A（形狀 10x10x36）和 B（形狀 10x27x36）。 我想乘以最后 2 個軸並沿軸 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支持堆疊。 你可以簡單地做：

 (A@B.transpose(0,2,1)).sum(0)

檢查（ C是使用 OP 的循環生成的）：

np.allclose((A@B.transpose(0,2,1)).sum(0),C)
# True
timeit(lambda:(A@B.transpose(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)