如何在 Python 中并行化循环?

How to parallelise loop in Python?

我正在寻找 运行 我的程序在本地多处理器上的瓶颈。我已经研究过 multiprocessing,但对于我的目的来说它看起来有点复杂,我想知道是否有更简单的方法。

我有一个超过 360 度角的循环,每个角的计算都是独立的,所以它以什么顺序完成并不重要。我有 8 个核心,所以希望我可以简单地将那个循环分成 8 个45 个角度的块并将它们发送到不同的核心并在最后收集结果。简化的示例如下所示:

dx = np.zeros(npixels)
for angle in range(360):
    dx += calculate_gradient_for_angle(angle, x, y, z, **kwargs)

这里,只有angle参数是可变的。其余为静态。

我查看了 multiprocessing.pool.Pool.map,但我只能找到显示传递给它的单参数函数的示例。如您所见,我的函数接受多个参数。任何指针将不胜感激。

我在 macOS 10.14.6

上使用 Python 3.7.8

你可以构建这样的东西。我的想法来自 this question

import concurrent.futures

def foo(bar):
    return bar

# Start each calculation in a seperate thread
with concurrent.futures.ThreadPoolExecutor() as ex:
    f = [ex.submit(foo, angle) for angle in range(10)]

dx = 0

# Collect all results
for r in f:
  dx += r.result()

print(dx)

简单的解决方案是使用 functools.partial 将函数转换为单参数函数并将其传递给 multiprocessing.pool.Pool.map,如下所示:

from functools import partial
import multiprocessing as mp

calculate_grad_partial = partial(calculate_grad_for_angle, x=x, y=y, z=z)

with mp.Pool(processes=8) as pool:
    dx = pool.map(calculate_grad_partial, range(360))