python - 如何在循环内使用 python 多处理 Pool.map

Question

我正在使用 Runge-Kutta 运行模拟。在每个时间步骤，两个自变量的两次 FFT 是必要的，可以并行化。我实现的代码是这样的:

from multiprocessing import Pool
import numpy as np

pool = Pool(processes=2)    # I like to calculate only 2 FFTs parallel 
                            # in every time step, therefor 2 processes

def Splitter(args):
    '''I have to pass 2 arguments'''
    return makeSomething(*args):

def makeSomething(a,b):
    '''dummy function instead of the one with the FFT'''
    return a*b

def RungeK():
    # ...
    # a lot of code which create the vectors A and B and calculates 
    # one Kunge-Kutta step for them 
    # ...

    n = 20                         # Just something for the example
    A = np.arange(50000)
    B = np.ones_like(A)

    for i in xrange(n):                  # loop over the time steps
        A *= np.mean(B)*B - A
        B *= np.sqrt(A)
        results = pool.map(Splitter,[(A,3),(B,2)])
        A = results[0]
        B = results[1]

    print np.mean(A)                                 # Some output
    print np.max(B)

if __name__== '__main__':
    RungeK()

不幸的是，Python 在到达循环后会生成无限数量的进程。之前看起来只有两个进程在运行。我的内存也被填满了。添加一个

pool.close()
pool.join()

循环后面并不能解决我的问题，将其放在循环内对我来说没有任何意义。希望您能帮忙。

Answer 1

将池的创建移至RungeK函数中；

def RungeK():
    # ...
    # a lot of code which create the vectors A and B and calculates
    # one Kunge-Kutta step for them
    # ...

    pool = Pool(processes=2)
    n = 20                         # Just something for the example
    A = np.arange(50000)
    B = np.ones_like(A)

    for i in xrange(n):  # loop over the time steps
        A *= np.mean(B)*B - A
        B *= np.sqrt(A)
        results = pool.map(Splitter, [(A, 3), (B, 2)])
        A = results[0]
        B = results[1]
    pool.close()
    print np.mean(A)  # Some output
    print np.max(B)

或者，将其放在主 block 中。

这可能是多处理工作方式的副作用。例如。在 MS Windows 上，您需要能够导入主模块而不产生副作用(例如创建新进程)。