在python多处理池中共享numpy数组

Question

我正在研究一些代码,这些代码在大量(数十到数十万个数值积分)问题上做了一些相当繁重的数值工作.幸运的是,这些集成非常平行,因此很容易使用Pool.map()将工作分成多个核心.

现在,我有一个具有这个基本工作流程的程序:

#!/usr/bin/env python
from multiprocessing import Pool
from scipy import *
from my_parser import parse_numpy_array
from my_project import heavy_computation

#X is a global multidimensional numpy array
X = parse_numpy_array("input.dat")
param_1 = 0.0168
param_2 = 1.505

def do_work(arg):
  return heavy_computation(X, param_1, param_2, arg)

if __name__=='__main__':
  pool = Pool()
  arglist = linspace(0.0,1.0,100)
  results = Pool.map(do_work,arglist)
  #save results in a .npy file for analysis
  save("Results", [X,results])

由于X,param_1和param_2是硬编码的,并且对于池中的每个进程以完全相同的方式进行初始化,因此一切正常.现在我的代码工作了,我想这样做,以便用户在运行时输入文件名,param_1和param_2,而不是硬编码.

应该注意的一件事是X,param_1和param_2在工作完成时没有被修改.由于我不修改它们,我可以在程序开始时做这样的事情:

import sys
X = parse_numpy_array(sys.argv[1])
param_1 = float(sys.argv[2])
param_2 = float(sys.argv[3])

这样做可以解决问题,但由于此代码的大多数用户都是从Windows机器运行代码,我宁愿不去命令行参数的路径.

我真正想做的是这样的事情:

X, param_1, param_2 = None, None, None

def init(x,p1, p2)
  X = x
  param_1 = p1
  param_2 = p2

if __name__=='__main__':
  filename = raw_input("Filename> ")
  param_1 = float(raw_input("Parameter 1: "))
  param_2 = float(raw_input("Parameter 2: "))
  X = parse_numpy_array(filename)
  pool = Pool(initializer = init, initargs = (X, param_1, param_2,))
  arglist = linspace(0.0,1.0,100)
  results = Pool.map(do_work,arglist)
  #save results in a .npy file for analysis
  save("Results", [X,results])

但是,当然,当pool.map调用发生时,这会失败并且X/param_1/param_2都是None.我对多处理很新,所以我不确定为什么对初始化程序的调用失败.有办法做我想做的事吗？有没有更好的方法来解决这个问题？我也看过使用共享数据,但是根据我对文档的理解,这只适用于ctypes,它不包括numpy数组.任何有关这方面的帮助将不胜感激.

Answer 1

我有一个类似的问题。如果您只想阅读我的解决方案，请跳过一些行:) 我必须：

• 在操作它的不同部分的线程之间共享一个 numpy.array 并且......
• 向 Pool.map 传递一个具有多个参数的函数。

我注意到：

• numpy.array 的数据被正确读取，但是......
• numpy.array 上的更改未永久化
• Pool.map 在处理 lambda 函数时遇到问题，或者在我看来是这样的（如果您不清楚这一点，请忽略它）

我的解决方案是：

• 使目标函数仅作为参数列表
• 使目标函数返回修改后的数据，而不是直接尝试写入 numpy.array

我知道您的 do_work 函数已经返回计算数据，因此您只需修改 to_work 以接受列表（包含 X、param_1、param_2 和 arg）作为参数，并在传递之前以这种格式将输入打包到目标函数它到 Pool.map。

这是一个示例实现：

def do_work2(args):
    X,param_1,param_2,arg = args
    return heavy_computation(X, param_1, param_2, arg)

现在您必须在调用 do_work 函数之前将输入打包。你的主要变成：

if __name__=='__main__':
   filename = raw_input("Filename> ")
   param_1 = float(raw_input("Parameter 1: "))
   param_2 = float(raw_input("Parameter 2: "))
   X = parse_numpy_array(filename)
   # now you pack the input arguments
   arglist = [[X,param1,param2,n] for n in linspace(0.0,1.0,100)]
   # consider that you're not making 100 copies of X here. You're just passing a reference to it
   results = Pool.map(do_work2,arglist)
   #save results in a .npy file for analysis
   save("Results", [X,results])