Python 多处理 - "outer product" 个参数

Python multiprocessing - "outer product" of arguments

我正在使用多个输入参数进行数值模拟。其中一些参数是静态的,而另一些是我想 运行 我的函数所在的数字数组。比如我可能要模拟下面这组参数

a = 1
b = np.arange(1, 11)
c = np.arange(20, 31)
d = 1

这意味着运行宁

simulate(a = 1, b = 1, c = 20, d = 1)
simulate(a = 1, b = 1, c = 21, d = 1)
...
simulate(a = 1, b = 1, c = 30, d = 1)
simulate(a = 1, b = 2, c = 20, d = 1)
...

simulate() 的 100 次调用。我想使用多线程来加快速度。我尝试使用 multiprocessing 的 pool.map(),但是输入参数的结构要求需要实例化一个长度为 100 的数组,其中包含 a、b、c 和 d 的列表,例如[[1, 1, 20, 1], [1, 1, 21, 1] ...

实际上我有足够的参数并且在足够的维度上变化,以至于我 运行 内存不足试图为 pool.map() 生成输入数组。

我想要的是有一个函数 map_wrapper() 这样

a = 1
b = np.arange(1, 11)
c = np.arange(20, 31)
d = 1
map_wrapper(simulation, [a, b, c, d])

相当于上面列出的simulation()的100次调用,或者是map()的一种使用方式或者类似的方式。

这里是一个最小的例子,展示了如何构建一个生成器来为多处理函数生成参数:

import multiprocessing as mp

def gen_args():
    for a in range(4):
        for b in range(4,8):
            for c in range(8,12):
                yield (a,b,c)
                
def foo(a, b, c):
    return a + b + c

if __name__ == '__main__':
    with mp.Pool() as p:
        res = p.starmap(foo, gen_args())

这将生成 64 个任务,但参数是根据需要即时计算的,而不是在开始时一次性全部计算。请记住,您仍然需要 space 作为输出列表。使用 starmap 的 chunksize 参数可能会也可能不会提高执行速度(通过两种方式进行测试以找出答案),但它会稍微增加内存使用量,因为它会为每个工作函数一次拉取多组参数继续努力。