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

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

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

Answer 2

此处使用 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 3

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

Numpy 3D 矩阵乘法

Numpy 3D matrix multiplication

python

numpy

matrix-multiplication