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 访问。

也可以考虑将数组分割成几块，然后使用线程（以避免复制数据）并在每块上并行应用函数。这比我正在寻找的更复杂（如果全局解释器锁没有被充分释放，可能无法正常工作）。

如果能够以简单的方式使用多个内核来执行简单的可并行任务，例如将函数应用于数组的所有元素（这基本上是这里需要的，函数 @ 987654327@ 采用一维参数数组。

Answer 1

好的，我解决了：一个想法是使用标准的multiprocessing 模块并将原始数组分成几块（以限制与工作人员的通信开销）。这可以相对容易地完成，如下所示：

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 本身）？