在 python 中并行化这些嵌套的 for 循环

Question

我有一个多维数组 (result)，应该由一些嵌套循环填充。函数fun() 是一个复杂且耗时的函数。我想以并行方式填充我的数组元素，这样我就可以使用我系统的所有处理能力。 代码如下：

import numpy as np


def fun(x, y, z):
    # time-consuming computation...
    # ...

    return output


dim1 = 10
dim2 = 20
dim3 = 30

result = np.zeros([dim1, dim2, dim3])

for i in xrange(dim1):
    for j in xrange(dim2):
        for k in xrange(dim3):
            result[i, j, k] = fun(i, j, k)

我的问题是“我可以并行化这段代码吗？如果可以，如何？”

我使用的是 Windows 10 64 位和 python 2.7。

如果可以的话，请通过更改我的代码来提供您的解决方案。 谢谢！

Answer 1

如果您想要一个更通用的解决方案，利用完全并行执行，那么为什么不使用这样的解决方案：

>>> import multiprocess as mp
>>> p = mp.Pool()
>>> 
>>> # a time consuming function taking x,y,z,...
>>> def fun(*args):
...   import time
...   time.sleep(.1)
...   return sum(*args)
... 
>>> dim1, dim2, dim3 = 10, 20, 30
>>> import itertools
>>> input = ((i,j,k) for i,j,k in itertools.combinations_with_replacement(xrange(dim3), 3) if i < dim1 and j < dim2)
>>> results = p.map(fun, input)
>>> p.close()
>>> p.join()
>>>
>>> results[:2]
[0, 1]
>>> results[-2:]
[56, 57]

请注意，我使用的是 multiprocess 而不是 multiprocessing，但这只是为了能够在解释器中工作。

我没有使用numpy.array，但如果你不得不...你可以直接将p.map 的输出转储到numpy.array，然后将shape 属性修改为shape = (dim1, dim2, dim3) ，或者你可以这样做：

>>> input = ((i,j,k) for i,j,k in itertools.combinations_with_replacement(xrange(dim3), 3) if i < dim1 and j < dim2)
>>> import numpy as np
>>> results = np.empty(dim1*dim2*dim3)
>>> res = p.imap(fun, input)
>>> for i,r in enumerate(res):
...   results[i] = r
... 
>>> results.shape = (dim1,dim2,dim3)

Answer 2

这是一个代码版本，它为不同的 k 索引并行运行 fun(i, j, k)。这是通过使用https://docs.python.org/2/library/multiprocessing.html在不同进程中运行fun来完成的

import numpy as np
from multiprocessing import Pool


def fun(x, y, z):
    # time-consuming computation...
    # ...

    return output


def fun_wrapper(indices):
    fun(*indices)

if __name__ == '__main__':
    dim1 = 10
    dim2 = 20
    dim3 = 30

    result = np.zeros([dim1, dim2, dim3])

    pool = Pool(processes=8)
    for i in xrange(dim1):
        for j in xrange(dim2):
            result[i, j] = pool.map(fun_wrapper, [(i, j, k) for k in xrange(dim3)])

这不是最优雅的解决方案，但您可以从它开始。只有当fun 包含耗时的计算时，您才能获得加速

Answer 3

一种简单的方法是将数组划分为多个部分并创建一些线程来通过这些部分进行操作。例如从 (0,0,0) 到 (5,10,15) 的一个部分，从 (5,10,16) 到 (10,20,30) 的另一个部分。

您可以使用threading 模块并执行类似的操作

import numpy as np
import threading as t


def fun(x, y, z):
    # time-consuming computation...
    # ...

    return output


dim1 = 10
dim2 = 20
dim3 = 30

result = np.zeros([dim1, dim2, dim3])
#b - beginning index, e - end index
def work(ib,jb,kb,ie,je,ke):
    for i in xrange(ib,ie):
        for j in xrange(jb,je):
            for k in xrange(kb,ke):
                result[i, j, k] = fun(i, j, k)

 threads = list()
 threads.append(t.Thread(target=work, args(0,0,0,dim1/2,dim2/2,dim3/2))
 threads.append(t.Thread(target=work, args(dim1/2,dim2/2,dim3/2 +1,dim1, dim2, dim3))

 for thread in threads:
     thread.start()

您可以通过某种算法定义这些部分并动态确定线程数。希望它可以帮助您或至少给您一些想法。