NumPy 中许多外部产品的总和答案

【问题标题】：Sum of many outer products in NumPyNumPy 中许多外部产品的总和
【发布时间】：2021-06-28 05:55:37
【问题描述】：

我有两个矩阵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 的解决方案可以解决问题，但我不熟悉它！

【问题讨论】：

【解决方案1】：

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

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

【讨论】：

【解决方案2】：

对于未来的读者：此操作在数学上相当于将A 的每一行的总和与B 的每一行的总和相乘。也就是说，最快的解决方案是C = np.sum(A,axis=-1) * np.sum(B,axis=-1)。

这比np.einsum 和任何其他计算方式要快得多。

【讨论】：