Python multiprocessing参数笛卡尔积并行调用的低内存实现咨询

核心思路是利用惰性迭代器替代预生成全量参数列表，全程不会一次性加载所有参数组合到内存，完美解决大参数规模下的内存不足问题。

依赖导入 #

import itertools
import multiprocessing as mp
import numpy as np
from collections.abc import Iterable

map_wrapper 实现 #

def map_wrapper(func, params, process_num=None, ordered=True):
    """
    惰性生成参数笛卡尔积，多进程调用目标函数
    :param func: 目标调用函数
    :param params: 参数列表，固定值和可迭代参数混合传入
    :param process_num: 进程数，默认取CPU核心数
    :param ordered: 是否按参数组合顺序返回结果，False时性能更高
    :return: 所有调用结果的列表
    """
    # 统一参数格式：固定值包装为单元素可迭代对象，排除字符串（避免被拆分为单个字符）
    formatted_params = []
    for p in params:
        if isinstance(p, str) or not isinstance(p, Iterable):
            formatted_params.append([p])
        else:
            formatted_params.append(p)
    
    # 生成惰性笛卡尔积迭代器，不会预先生成全量组合
    param_iter = itertools.product(*formatted_params)
    
    # 多进程执行，自动解包参数
    def args_unpack_wrapper(args):
        return func(*args)
    
    with mp.Pool(processes=process_num) as pool:
        if ordered:
            res = list(pool.imap(args_unpack_wrapper, param_iter))
        else:
            res = list(pool.imap_unordered(args_unpack_wrapper, param_iter))
    
    return res

使用示例 #

# 你的模拟函数示例
def simulate(a, b, c, d):
    # 这里写你的模拟逻辑
    return a + b + c + d

if __name__ == '__main__':
    # 按你的方式定义参数即可，固定值和数组混合
    a = 1
    b = np.arange(1, 11)
    c = np.arange(20, 31)
    d = 1
    
    # 调用方法和你要求的完全一致
    results = map_wrapper(simulate, [a, b, c, d])
    # 可选非顺序返回提升性能：results = map_wrapper(simulate, [a, b, c, d], ordered=False)
    # 可选指定进程数：results = map_wrapper(simulate, [a, b, c, d], process_num=8)

关键说明 #

• itertools.product生成的是惰性迭代器，仅当多进程拉取任务时才会生成下一组参数，哪怕参数组合有百万、千万级也不会占用过量内存
• 若不需要按参数顺序返回结果，设置ordered=False可以获得更高的执行效率
• Windows系统下必须将调用逻辑放在if __name__ == '__main__'代码块中，否则会出现多进程启动异常

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