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

Question

我有一个复杂的numpy数组signal，其维度为[10,1000,50000]。 我需要按切片修改这个数组。这是通过for循环完成的：

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

我尽可能地优化了myfunction。当我运行脚本时，它需要相当长的时间，但只使用24个CPU中的1个。

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

因此，我想使用并行计算来加速我的代码。 在Python中，似乎有许多不同的并行计算方法。 哪种方法似乎最适合我的问题？我该如何实现它？

- torpedo

---

一般来说，当你在使用NumPy时，你要尽量避免在Python代码中使用循环，而是要找到正确的NumPy函数，这些函数将能够一次性处理整个数组。然后就交给NumPy去处理并行化了。 - zvone

myfunction是什么？ - Daniel

如果函数的参数是从切片计算出来的，那么你可以告诉numpy为每个切片执行这个操作。 - zvone

好的，大家好。首先感谢所有的反馈。我无法对整个数组执行计算以检索每个切片的参数。我需要在每个切片上独立执行myfunction。myfunction太复杂了。再次感谢。 - torpedo

使用 Dask。它被设计用于并行化 NumPy 操作。 - Imanol Luengo

Answer 1

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

- Deena

---

代码运行得很好。但是我也尝试将其应用于一个4D数组，通过访问np_array[:,i,:,:]。然而，在其中一个joblib函数中，我收到了错误消息“输出数组是只读的”。 - torpedo