并行修改切片中的3D numpy数组

时间:2018-09-01 17:07:40

标签: python parallel-processing multiprocessing

我有一个复杂的numpy数组 signal ,其尺寸为[10,1000,50000] 我需要在切片中修改此数组。这是在for循环中完成的:

for k in range(signal.shape[2]):
    signal[:,:,k] = myfunction(signal[:,:,k], constant1, constant2, constant5=constant5, constant6=constant6)

我尽可能地优化了我的功能。当我运行脚本时,需要花费一些时间,但只使用24个CPU中的1个。

无法使用numpy重写代码以在整个数组上执行myfunction。

因此,我想通过并行计算来加速我的代码。 在python中似乎有许多不同的并行计算方法。 哪一个似乎最适合我的问题?我该如何实施呢?

1 个答案:

答案 0 :(得分:1)

Joblib为此类“令人尴尬的并行”任务提供了轻松执行:

import numpy as np

# Initialize array and define function 
np_array = np.random.rand(100,100,100)
my_function = lambda x: x / np.sum(x)

# Option 1: Loop over array and apply function
serial_result = np_array.copy()
for i in range(np_array.shape[2]):
    serial_result[:,:,i] = my_function(np_array[:,:,i])

现在对joblib使用并行执行:

# Option 2: Parallel execution
# ... Apply function in Parallel 
from joblib import delayed, parallel
sub_arrays = Parallel(n_jobs=6)(                            # Use 6 cores 
                      delayed(my_function)(np_array[:,:,i]) # Apply my_function 
                      for i in range(np_array.shape[2]))    # For each 3rd dimension

# ... Concatenate the list of returned arrays
parallel_results = np.stack(sub_arrays, axis=2)

# Compare results 
np.equal(serial_result, parallel_results).all() # True