Распараллеливать эти вложенные петли в python

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-bit и python 2.7.

Пожалуйста, предоставьте свое решение, изменив свой код, если сможете. Спасибо!

Answer 1

Простым подходом может быть разделение массива на разделы и создание потоков для этих разделов. Например, один участок от (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() 

Вы можете определить эти разделы с помощью некоторого алгоритма и определить количество потоков динамически. Надеюсь, это поможет вам или, по крайней мере, дать вам некоторые идеи.

Answer 2

Ниже приведена версия кода, которая работает fun(i, j, k) в качестве альтернативы для differend k индексов. Это делается путем запуска fun в различных процессах с использованием https://docs.python.org/2/library/multiprocessing.html

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

Если вы хотите получить более общее решение, воспользовавшись полностью параллельного выполнения, то почему бы не использовать что-то вроде этого:

>>> 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)