python - numpy.apply_along_axis() 的简单并行化？

Question

如何通过 numpy.apply_along_axis() 将函数应用到 NumPy 数组的元素并进行并行化，以便利用多核？这似乎是一件自然而然的事情，在对所应用函数的所有调用都是独立的常见情况下。

在我的特殊情况下——如果这很重要的话——应用轴是轴 0:np.apply_along_axis(func, axis=0, arr=param_grid) (np 成为 NumPy)。

我快速查看了 Numba，但我似乎无法通过如下循环实现这种并行化:

@numba.jit(parallel=True)
result = np.empty(shape=params.shape[1:])
for index in np.ndindex(*result.shape)):  # All the indices of params[0,...]
    result[index] = func(params[(slice(None),) + index])  # Applying func along axis 0

显然还有一个 NumPy 中的编译选项 用于通过 OpenMP 进行并行化，但它似乎无法通过 MacPorts 访问。

也可以考虑将数组分成几部分并使用线程(以避免复制数据)并在每个部分上并行应用该函数。这比我正在寻找的更复杂(如果 Global Interpreter Lock 未充分释放，则可能无法工作)。

如果能够以一种简单的方式使用多个内核来执行简单的可并行任务，比如将一个函数应用于数组的所有元素(这实际上是这里所需要的，但函数 func() 采用一维参数数组)。

Answer 1

好吧，我想出来了:一个想法是使用标准的 multiprocessing 模块并将原始数组拆分为几个 block (以限制与工作人员的通信开销)。这可以相对容易地完成，如下所示:

import multiprocessing

import numpy as np

def parallel_apply_along_axis(func1d, axis, arr, *args, **kwargs):
    """
    Like numpy.apply_along_axis(), but takes advantage of multiple
    cores.
    """        
    # Effective axis where apply_along_axis() will be applied by each
    # worker (any non-zero axis number would work, so as to allow the use
    # of `np.array_split()`, which is only done on axis 0):
    effective_axis = 1 if axis == 0 else axis
    if effective_axis != axis:
        arr = arr.swapaxes(axis, effective_axis)

    # Chunks for the mapping (only a few chunks):
    chunks = [(func1d, effective_axis, sub_arr, args, kwargs)
              for sub_arr in np.array_split(arr, multiprocessing.cpu_count())]

    pool = multiprocessing.Pool()
    individual_results = pool.map(unpacking_apply_along_axis, chunks)
    # Freeing the workers:
    pool.close()
    pool.join()

    return np.concatenate(individual_results)

在 Pool.map() 中应用的函数 unpacking_apply_along_axis() 是独立的(以便子进程可以导入它)，并且只是一个薄的处理 Pool.map() 只接受一个参数这一事实的包装器:

def unpacking_apply_along_axis((func1d, axis, arr, args, kwargs)):
    """
    Like numpy.apply_along_axis(), but with arguments in a tuple
    instead.

    This function is useful with multiprocessing.Pool().map(): (1)
    map() only handles functions that take a single argument, and (2)
    this function can generally be imported from a module, as required
    by map().
    """
    return np.apply_along_axis(func1d, axis, arr, *args, **kwargs)

(在 Python 3 中，这应该写成

def unpacking_apply_along_axis(all_args):
    (func1d, axis, arr, args, kwargs) = all_args

因为argument unpacking was removed ).

在我的特殊情况下，这导致 2 个内核的超线程速度提高了 2 倍。接近 4 倍的因子会更好，但速度已经很好了，只需几行代码，对于具有更多内核(这很常见)的机器来说应该更好。也许有一种避免数据复制和使用共享内存的方法(可能通过 multiprocessing module 本身)？