大型numpy数组执行聚合运算内存溢出 单机器环境下如何解决

核心思路 #

计算axis=0维度的均值，本质是对生成器返回的100个shape为(10**6, 133)的数组做对应位置的求和后除以数组总数。不需要一次性加载所有数组到内存，采用增量计算即可将内存占用控制在单个数组的2倍以内，完全适配单台普通配置机器的运行要求。

实现代码 #

方案1：累加求和后取均值（性能最优） #

该方案实现简单，计算效率最高，适合数值范围不会触发累加溢出的场景：

import numpy as np

# 你的业务数据生成器
data = (np.ones((10**6, 133)) for _ in range(100))

sum_array = None
count = 0
for arr in data:
    if sum_array is None:
        # 初始化累加数组，用float64避免精度损失
        sum_array = np.zeros_like(arr, dtype=np.float64)
    sum_array += arr
    count += 1

# 最终计算均值
mean_array = sum_array / count

该方案内存峰值仅为1GB左右（累加数组约532MB+单批次读取数组约532MB），无额外磁盘IO开销。

方案2：增量更新均值（避免溢出） #

如果单个数组数值极大，担心累加过程出现溢出，可以采用增量更新均值的方式，全程不需要存储总和：

import numpy as np

data = (np.ones((10**6, 133)) for _ in range(100))

mean_array = None
count = 0
for arr in data:
    count += 1
    if mean_array is None:
        mean_array = arr.astype(np.float64)
    else:
        # 增量更新公式，避免大数累加溢出
        mean_array += (arr - mean_array) / count

原方法报错原因 #

• 直接用np.stack会一次性把生成器返回的所有100个数组加载到内存，需要约53GB内存，远超普通机器配置触发内存错误。
• 你采用的dask写法不会自动将内存中的numpy数组落盘，构建dask数组时依然会全量加载所有数据到内存，触发系统OOM机制直接杀死进程。如果确实需要用dask实现，需要先将所有数组落地为磁盘npy文件再读取分块计算，复杂度远高于增量计算方案。

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