两个for循环的多重处理

Question

我正在努力在 python (2.7) 中实现算法以并行化物理问题的计算。在两个变量（假设 a 和 b）上有一个参数空间，我想在上面运行我编写的程序 f(a,b)，它返回另外两个变量 c 和 d。到目前为止，我for在 a 和 b 上使用了两个循环来计算 c 和 d 的两个数组，然后将它们保存为 txt 文档。由于参数空间比较大，每次计算其中的一个点f(a,b)的时间都比较长，如果能把我的8个CPU核全部用于参数空间扫描就好了。我已经阅读了有关多线程和多处理的信息，似乎多处理就是我正在寻找的东西。你知道这个应用程序的一个很好的代码示例或资源来了解我相当简单的应用程序的多处理基础知识吗？

Answer 1

下面是一个示例，说明如何将 multiprocessing 与一个简单函数一起使用，该函数接受两个参数并返回一个包含两个数字的元组，以及一个您要在其上进行计算的参数空间：

from itertools import product

from multiprocessing import Pool

import numpy as np

def f(a, b):

    c = a + b

    d = a * b

    return (c, d)

a_vals = [1, 2, 3, 4, 5, 6]

b_vals = [10, 11, 12, 13, 14, 15, 16, 17]

na = len(a_vals)

nb = len(b_vals)

p = Pool(8)  # <== maximum number of simultaneous worker processes

answers = np.array(p.starmap(f, product(a_vals, b_vals))).reshape(na, nb, 2)

c_vals = answers[:,:,0]

d_vals = answers[:,:,1]

这给出了以下内容：

>>> c_vals

array([[11, 12, 13, 14, 15, 16, 17, 18],

       [12, 13, 14, 15, 16, 17, 18, 19],

       [13, 14, 15, 16, 17, 18, 19, 20],

       [14, 15, 16, 17, 18, 19, 20, 21],

       [15, 16, 17, 18, 19, 20, 21, 22],

       [16, 17, 18, 19, 20, 21, 22, 23]])

>>> d_vals

array([[ 10,  11,  12,  13,  14,  15,  16,  17],

       [ 20,  22,  24,  26,  28,  30,  32,  34],

       [ 30,  33,  36,  39,  42,  45,  48,  51],

       [ 40,  44,  48,  52,  56,  60,  64,  68],

       [ 50,  55,  60,  65,  70,  75,  80,  85],

       [ 60,  66,  72,  78,  84,  90,  96, 102]])

返回p.starmap一个二元组列表，然后从中提取 c 和 d 值。

这假定您将在取回所有结果后在主程序中执行文件 I/O。

附录：

如果p.starmap不可用（Python 2），那么您可以更改函数以采用单个输入（2 元素元组）：

def f(inputs):

    a, b = inputs

    # ... etc as before ...

然后在上面的代码中使用p.map代替。p.starmap

如果不方便更改函数（例如它也从其他地方调用），那么您当然可以编写一个包装函数：

def f_wrap(inputs):

    a, b = inputs

    return f(a, b)

并改为调用它。