Python Numpy三维矩阵乘法

我有两个矩阵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)

---

我想以矢量化的方式实现这一点，但似乎不知道如何实现。我已经检查过了，但还不确定如何应用它来实现结果。任何帮助都将不胜感激。谢谢

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

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

---

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

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

---

这里使用

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)

---

这里使用

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)

---

matmul

支持堆叠。您可以简单地执行以下操作：

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

检查（

C

使用OP的循环生成）：

---

matmul

支持堆叠。您可以简单地执行以下操作：

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

检查（

C

使用OP的循环生成）：

---

不幸的是，它仍然涉及一个本机Python列表，因此与我的原始代码的唯一区别是syntaxGreat解决方案。不幸的是，它仍然涉及一个本机Python列表，因此与我的原始代码的唯一区别也是syntaxGreat解决方案。谢谢这也是一个很好的解决方案。谢谢