在NumPy中计算多个外积的总和

Question

我有两个矩阵A和B，都具有形状(N,M)。我想执行以下操作：C = np.sum(A[:,None,:]*B[:,:,None],axis=(1,2))，它对应于对每行的A与每行的B进行外积求和。然后，C的形状将为(N,)。

问题在于，当使用这种形式时，我会得到一个MemoryError，因为N=12000且M=4000。

有没有一种方法可以执行此操作，而无需先构建要求和的（巨大的）中间数组？

我怀疑使用np.einsum的解决方案会奏效，但我对其不熟悉！

- Toool

Answer 1

不确定np.einsum是否解决了内存问题，但以下是使用它计算相当于您的计算：

C = np.einsum('ij,ik->i',A,B)

- Ehsan

---

谢谢！np.einsum 明显更快，使用的内存也更少，我会检查它是否解决了我的问题。 - Toool

Answer 2

针对未来读者：这个操作在数学上等同于将A的每一行之和与B的每一行之和相乘。换句话说，最快的解决方案是 C = np.sum(A,axis=-1) * np.sum(B,axis=-1)。

这比使用np.einsum和其他任何计算方法都要快得多。

- Toool