python - 通过多处理对生成器输出应用 reduce

Question

我有一个像这样工作的生成器函数 (Python)

def Mygenerator(x, y, z, ...):
    while True:
        # code that makes two matrices based on sequences of input arrays
        yield (matrix1, matrix2)

我想做的是添加这个生成器的输出。这条线完成了这项工作:

M1, M2 = reduce(lambda x, y: x[0] + y[0], x[1] + y[1], Mygenerator(x, y, z, ...))

我想将其并行化以加快计算速度。重要的是，Mygenerator 的输出在生成时会减少，因为 list(Mygenerator(...)) 会占用太多内存。

Answer 1

为了回答我自己的问题，我找到了一个似乎像我希望的那样工作的解决方案:

首先，Mygenerator 不再是一个生成器，而是一个函数。此外，我现在不再循环遍历 x、y 和 z 的片段，而是将一个片段传递给当时的函数:

def Myfunction(x_segment, y_segment, z_segment):
        # code that makes two matrices based on input arrays
        return (matrix1, matrix2)

将 multiprocessing.Pool 与 imap(生成器)函数一起使用似乎可行:

pool = multiprocessing.Pool(ncpus)
results = pool.imap(Myfunction, 
                    ( (x[i], y[i], z[i]) for i in range(len(x)) )
M1, M2 = reduce(lambda r1, r2: (r1[0] + r2[0], r1[1] + r2[1]), 
                    (result for result in results))
pool.close()
pool.join()

我将 lambda 表达式中的 x 和 y 更改为 r1 和 r2 以避免混淆具有相同名称的其他变量。当尝试将生成器与 multiprocessing 一起使用时，我在 pickle 上遇到了一些麻烦。

这个解决方案唯一令人失望的是它并没有真正加快计算速度。我想这与开销操作有关。使用 8 核时，处理速度提高了约 10%。当减少到 4 个核心时，速度翻了一番。这似乎是我对我的特定任务所能做的最好的，除非有一些其他的方法来进行并行化......

这里必须使用imap 函数，因为map 会在reduce 操作之前将所有返回值存储在内存中，而在这种情况是不可能的。