python - 如何在 python numpy 中并行化求和计算？

Question

我有一个要计算的总和，但我很难并行化代码。我尝试并行化的计算有点复杂(它同时使用 numpy 数组和 scipy 稀疏矩阵)。它吐出一个 numpy 数组，我想对大约 1000 次计算的输出数组求和。理想情况下，我会在所有迭代中保持运行总和。但是，我一直无法弄清楚如何执行此操作。

到目前为止，我已经尝试将 joblib 的 Parallel 函数和 pool.map 函数与 python 的多处理包一起使用。对于这两者，我使用了一个返回 numpy 数组的内部函数。这些函数返回一个列表，我将其转换为一个 numpy 数组，然后求和。

然而，joblib Parallel 函数完成所有迭代后，主程序再也没有继续运行(看起来原来的作业处于挂起状态，使用 0% CPU)。当我使用 pool.map 时，在所有迭代完成后出现内存错误。

有没有办法简单地并行化数组的运行总和？

编辑:目标是执行类似以下的操作，但并行操作除外。

def summers(num_iters):

    sumArr = np.zeros((1,512*512)) #initialize sum
    for index in range(num_iters):
        sumArr = sumArr + computation(index) #computation returns a 1 x 512^2 numpy array

    return sumArr

Answer 1

我弄清楚了如何使用多处理、apply_async 和回调对数组总和进行并行处理，所以我将其发布在这里供其他人使用。我用了the example page for Parallel Python对于 Sum 回调类，虽然我实际上并没有使用那个包来实现。不过，它让我想到了使用回调。这是我最终使用的简化代码，它完成了我想要它做的事情。

import multiprocessing
import numpy as np
import thread

class Sum: #again, this class is from ParallelPython's example code (I modified for an array and added comments)
    def __init__(self):
        self.value = np.zeros((1,512*512)) #this is the initialization of the sum
        self.lock = thread.allocate_lock()
        self.count = 0

    def add(self,value):
        self.count += 1
        self.lock.acquire() #lock so sum is correct if two processes return at same time
        self.value += value #the actual summation
        self.lock.release()

def computation(index):
    array1 = np.ones((1,512*512))*index #this is where the array-returning computation goes
    return array1

def summers(num_iters):
    pool = multiprocessing.Pool(processes=8)

    sumArr = Sum() #create an instance of callback class and zero the sum
    for index in range(num_iters):
        singlepoolresult = pool.apply_async(computation,(index,),callback=sumArr.add)

    pool.close()
    pool.join() #waits for all the processes to finish

    return sumArr.value

我还能够使用并行映射来实现这一点，这是在另一个答案中提出的建议。我之前试过这个，但我没有正确实现。两种方法都有效，我认为 this answer很好地解释了使用哪种方法(map 或 apply.async)的问题。对于 map 版本，您不需要定义类 Sum 并且 summers 函数变为

def summers(num_iters):
    pool = multiprocessing.Pool(processes=8)

    outputArr = np.zeros((num_iters,1,512*512)) #you wouldn't have to initialize these
    sumArr = np.zeros((1,512*512))              #but I do to make sure I have the memory

    outputArr = np.array(pool.map(computation, range(num_iters)))
    sumArr = outputArr.sum(0)

    pool.close() #not sure if this is still needed since map waits for all iterations

    return sumArr