Python多处理用于2D数组的昂贵操作

我有一个非常昂贵的函数，比如

fun（a，b）

，它返回一组数据

我当前的方法如下所示：

a = np.linspace(0,100,300)
b = np.linspace(0,100,300)
A,B = np.meshgrid(a,b)
Y = np.zeros(A.shape)

for i,av in enumerate(a):
  for j, bv in enumerate(b):
    Y[i,j] = fun(av,bv)

---

（在某种程度上，我不得不把事情弄得一团糟，以使它适合我）。无论如何，这个过程需要相当长的时间，我想知道是否有一种简单的方法可以使用我的多核处理器来加快速度

有一个很棒的模块，名为python，它是python标准库的一部分。它将在尽可能多的内核上生成进程，以利用其他CPU。这里有一个在文档中使用对象的示例，下面是该示例的简化版本。它将计算10个数字的平方，将工作负载分布到各个工作流程中，并显示结果

简单工作人员池

from multiprocessing import Pool

def f(x):
    return x*x

pool = Pool(processes=4)
print pool.map(f, range(10))

输出

[0, 1, 4, 9, 16, 25, 36, 49, 64, 81]

{(0, 0): (100, 100),
 (0, 1): (100, 200),
 (0, 2): (100, 300),
 (1, 0): (200, 100),
 (1, 1): (200, 200),
 (1, 2): (200, 300),
 (2, 0): (300, 100),
 (2, 1): (300, 200),
 (2, 2): (300, 300)}

我面临的挑战更多的是将你的问题分解成相同的结构。我必须创建一些中介功能来完成这项工作。我没有numpy，所以我只是用列表和字典来代替你放的东西。您可以替换它们并尝试代码

更复杂的场景

from multiprocessing import Pool
import time, pprint

def fun(av, bv):
    time.sleep(0.1)
    return (av, bv)

def data_stream(a, b):
    for i, av in enumerate(a):
        for j, bv in enumerate(b):
            yield (i, j), (av, bv)

def proxy(args):
    return args[0], fun(*args[1])

a = range(100, 400, 100)
b = range(100, 400, 100)
Y = {}

pool = Pool(processes=4)
results = pool.map(proxy, data_stream(a, b))
for k,v in results:
    Y[k] = v

pprint.pprint(Y)

输出

[0, 1, 4, 9, 16, 25, 36, 49, 64, 81]

{(0, 0): (100, 100),
 (0, 1): (100, 200),
 (0, 2): (100, 300),
 (1, 0): (200, 100),
 (1, 1): (200, 200),
 (1, 2): (200, 300),
 (2, 0): (300, 100),
 (2, 1): (300, 200),
 (2, 2): (300, 300)}

性能

在这个例子中，我只是放置了一个虚拟的0.1秒延迟来模拟繁重的工作。但是，即使在本例中，如果您使用

processs=1

运行一个池，它也会使用

processs=4

以0.352s的速度运行。您可以以多种不同的方式使用多处理库，池只是一种方式。您可能希望探索示例和实验

---

在下面的一条评论中，提到了对pool.map使用chunksize参数来帮助提高性能。对引擎盖下发生的事情有一个大致的了解对于真正掌握性能是很重要的。基本上，传递给其他进程的所有数据都需要经过pickle处理，然后传递给未使用pickle的其他进程，然后结果通过相同的进程返回到主进程。这种进程间通信存在开销。实验时请记住这一点。

也许这个问题（及其答案）会有所帮助：在旁注中，请看一下

numpy.ndenumerate

。它和嵌套for循环之间没有性能差异，但是它更干净（对于

pool.imap

等东西来说非常方便）。除非您想在更深入的并行化方面投入时间，否则这似乎是最好的选择。您还可以提到将

chunksize

参数增加到

pool.map

（或

pool.imap

）。它通常会显著提高执行速度，特别是对于访问单个项的开销相对较高的numpy数组这样的情况。