使用multiprocessing.map处理多进程任务时如何节省内存避免超限？

核心思路是放弃存储全部100万个中间向量，仅保留最终累加结果，通过增量计算或分块处理将内存占用控制在可控范围，可选方案如下：

方案1：增量迭代累加（改造成本最低） #

将原来的pool.map替换为imap_unordered，获取结果迭代器逐个累加，无需等待所有计算完成也无需全量存储中间结果，内存仅占用1个最终累加向量+当前批次的少量中间向量：

import numpy as np
import multiprocessing

# 你原有的函数，输入(x,y)元组，返回长度3000的向量
def f(xy_pair):
    x, y = xy_pair
    # 原有逻辑不变
    pass

if __name__ == "__main__":
    # 你原有的x、y笛卡尔积集合
    RANGE = ...
    # 初始化累加结果，数据类型按需调整精度
    sum_result = np.zeros(3000, dtype=np.float64)
    
    with multiprocessing.Pool() as pool_obj:
        # chunksize可根据函数计算速度调整，计算越快chunksize可以设越大，减少进程通信开销
        for vec in pool_obj.imap_unordered(f, RANGE, chunksize=200):
            sum_result += vec

方案2：分块批量处理（适合函数计算速度快的场景） #

如果单个f的计算耗时很低，逐个累加的进程通信开销会很高，可以将总任务拆分为多个块，每个块内先完成局部累加再返回主进程合并，大幅减少通信次数：

import numpy as np
import multiprocessing

def f(xy_pair):
    x, y = xy_pair
    # 原有逻辑不变
    pass

if __name__ == "__main__":
    RANGE = ...
    # 每块包含1000个(x,y)对，总共有1000个块，可按需调整块大小
    chunk_size = 1000
    chunks = [RANGE[i*chunk_size : (i+1)*chunk_size] for i in range(len(RANGE)//chunk_size)]
    
    # 子进程处理单个块，返回块内累加结果
    def process_chunk(chunk):
        chunk_sum = np.zeros(3000, dtype=np.float64)
        for xy in chunk:
            chunk_sum += f(xy)
        return chunk_sum
    
    sum_result = np.zeros(3000, dtype=np.float64)
    with multiprocessing.Pool() as pool:
        for chunk_sum in pool.imap_unordered(process_chunk, chunks):
            sum_result += chunk_sum

可选优化点 #

• 若对精度要求不高，将向量数据类型从float64改为float32，内存占用和计算速度均可优化50%左右
• 如果你的函数f满足可拆分特性，即f(x,y) = f_x(x) + f_y(y)，可直接分别计算所有x、y对应向量的和，总结果为sum_fx * len(y_list) + sum_fy * len(x_list)，计算复杂度直接从1e6降至2e3，效率提升最明显

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