Python多处理,用于2D阵列的昂贵操作

Question

我有一个功能,说它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)

(排序,我不得不混淆一些东西以使其适合).无论如何,这个过程需要相当长的时间,我想知道是否有一种直接的方式来使用我的多核处理器来加快速度.

Answer 1

有一个很棒的模块叫做multiprocessing,它是python标准库的一部分.它会在您希望利用其他CPU的多个内核中生成进程.有一个在文档中使用Pool对象的示例,下面是该示例的缩短版本.它将计算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]

我有更多的挑战将你的问题分解成同样的结构.我不得不创建一些中间函数来实现这项工作.我没有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, 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秒的延迟来模拟繁重的工作.但即使在这个例子中,如果你运行一个池,processes=1它运行在0.950s,processes=4它运行在0.352s.您可以通过多种方式使用多处理库,Pool只是一种方式.您可能想要探索示例和实验.

在下面的一条评论中,提到了使用pool.map的chunksize参数来帮助提高性能.重要的是要全面了解底层的内容,以便真正掌握性能.基本上,您传递给其他进程的所有数据都需要被pickle传递给其他未使用的进程,然后结果将通过相同的进程返回到主进程.这种进程间通信存在开销.在实验时请记住这一点.