python - 如何使用多处理将函数应用于二维 numpy 数组

Question

假设我有以下功能:

def f(x,y):
    return x*y

如何使用多处理模块将函数应用于 NxM 二维 numpy 数组中的每个元素？使用串行迭代，代码可能如下所示:

import numpy as np
N = 10
M = 12
results = np.zeros(shape=(N,M))
for x in range(N):
    for y in range(M):
        results[x,y] = f(x,y)

Answer 1

下面是您可以如何使用 multiprocessing 并行化您的示例函数。我还包含了一个几乎完全相同的纯 Python 函数，它使用非并行 for 循环，以及一个实现相同结果的 numpy 单行代码:

import numpy as np
from multiprocessing import Pool


def f(x,y):
    return x * y

# this helper function is needed because map() can only be used for functions
# that take a single argument (see http://stackoverflow.com/q/5442910/1461210)
def splat_f(args):
    return f(*args)

# a pool of 8 worker processes
pool = Pool(8)

def parallel(M, N):
    results = pool.map(splat_f, ((i, j) for i in range(M) for j in range(N)))
    return np.array(results).reshape(M, N)

def nonparallel(M, N):
    out = np.zeros((M, N), np.int)
    for i in range(M):
        for j in range(N):
            out[i, j] = f(i, j)
    return out

def broadcast(M, N):
    return np.prod(np.ogrid[:M, :N])

现在让我们看看性能:

%timeit parallel(1000, 1000)
# 1 loops, best of 3: 1.67 s per loop

%timeit nonparallel(1000, 1000)
# 1 loops, best of 3: 395 ms per loop

%timeit broadcast(1000, 1000)
# 100 loops, best of 3: 2 ms per loop

非并行的纯 Python 版本比并行版本高出大约 4 倍，而使用 numpy 数组广播的版本绝对碾压另外两个。

问题在于启动和停止 Python 子进程会带来大量开销，而您的测试函数又是如此微不足道，以至于每个工作线程只花费其生命周期的一小部分时间来做有用的工作。只有当每个线程在被杀死之前都有大量工作要做时，多处理才有意义。例如，您可以给每个工作人员一个更大的输出数组 block 来计算(尝试将 chunksize= 参数弄乱给 pool.map())，但是对于这样一个微不足道的例子，我怀疑你会看到很大的改进。

我不知道您的实际代码是什么样的 - 也许您的函数足够大且昂贵，足以保证使用多处理。但是，我敢打赌有很多更好的方法来提高其性能。