Numpy中如何高效实现外积加权求和（规避内存溢出）

你要计算的是加权外积的累加和，完全可以通过调整NumPy的向量化操作避免生成超大中间数组，同时保证运算效率：

方法1：优化Einstein求和表达式 #

直接将三个数组传入np.einsum，让它在内部完成元素级加权和累加，不会生成(60000,784,784)的中间数组：

import numpy as np
S = np.einsum('i,ij,ik->jk', A, B, B)

表达式解释：

• i对应数组A的维度，ij对应B的每行（60000个784维向量），ik对应B的每行转置视角
• 对i维度求和，最终得到784×784的结果矩阵，内存占用仅保留最终结果和临时计算的小批量数据

方法2：利用矩阵乘法（BLAS加速） #

该运算可以转化为矩阵乘法形式，NumPy会自动调用优化后的BLAS库，速度通常更快：

S = B.T @ (A[:, np.newaxis] * B)

原理说明：

• A[:, np.newaxis] * B 会通过广播生成(60000,784)的数组，每个元素是A[i] * B[i,j]
• 再和B.T（形状784×60000）做矩阵乘法，本质就是对每个i计算B[:,i] * A[i] * B[i,:]的累加，最终得到目标矩阵

这两种方法的内存占用都远低于原始方案，且运算速度和纯向量化操作一致，比Python循环快几个数量级。

内容的提问来源于stack exchange，提问作者user007