转换嵌套循环 Python 代码以使用多处理池

Converting nested loop Python code to use multiprocessing pool

我有一些代码具有以下结构:一个接受嵌套参数范围的函数。对于 MVCE 目的,考虑以下内容就足够了

def func(x, y):
    return x*x, x+(5*y)

results = []
for a in range(3):
    for b in range(2,5):
        temp_results = func(a,b)
        results.append(temp_results)

sorted_results = sorted(results, key=lambda x: x[1], reverse=True)
print(sorted_results)

我正在尝试将多处理池用于 运行 func(a, b) 部分。我想尽量保持上面的结构。

以下是目前为止我所看到的基于多处理池示例的内容。地图参数不起作用,因为他们期望其中的块大小。另外,如何将池结果收集到列表中以便稍后排序?有没有办法以与上述单线程情况非常相似的方式来做到这一点?

a_params = []
b_params = []
for a in range(3):
    for b in range(2,5):
        a_params.append(a)
        b_params.append(b)

from multiprocessing import Pool

def pool_handler():
    p = Pool(2)
    p.map(func, a_params, b_params)

if __name__ == '__main__':
    pool_handler()

map 适用于具有单个参数的函数。如果要传递多个参数,请使用 starmap(或重写 func 以获取参数元组并将参数元组传递给 map)。

您无需执行任何额外操作即可将结果放入列表中。 map/starmap自动将结果累积到列表中。

所以,总的来说,这是可行的:

import multiprocessing as mp

def func(x, y):
    print(x*x, x+(5*y))
    return x*x, x+(5*y)

def main():
    with mp.Pool(2) as pool:
        results = pool.starmap(func,
                               [(a, b) for a in range(3)
                                       for b in range(2, 5)])
    sorted_results = sorted(results, key=lambda x: x[1], reverse=True)
    print(sorted_results)

if __name__ == '__main__':
    main()
    

并显示:

0 10
0 15
0 20
1 11
1 16
1 21
4 12
4 17
4 22
[(4, 22), (1, 21), (0, 20), (4, 17), (1, 16), (0, 15), (4, 12), (1, 11), (0, 10)]

您只需要保存结果,最后下单即可。

import multiprocessing as mp


def func(vals):
    x, y = vals
    return x*x, x+(5*y)


def main():
    input_data = []
    for a in range(3):
        for b in range(2, 5):
            input_data.append([a, b])

    n_threads = 4
    pool = mp.Pool
    with pool(n_threads) as p:
        results = p.map(func, input_data)

    sorted_results = sorted(results, key=lambda x: x[1], reverse=True)
    print(sorted_results)


if __name__ == '__main__':
    main()