以节省时间和内存的方式使用多处理

Question

我有一个问题，我的想法已经不多了，所以我希望有人有更好的想法。

我正在尝试在 python 中对数百万个数据点进行优化。我有一个函数calculate(data)，它接收一个巨大的数组data，并返回一个大小相同的数组results。计算相当简单，但需要使用data 中的多个条目对results 中的每个条目进行计算，不幸的是无法矢量化。由于维度原因，data 和 results 的大小很大，无法缩小。 由于计算的数量庞大，这可能需要几天的时间来计算。

因此我开始使用multiprocessing，它显着提高了速度。我实现这一点的方法是将results 切割成块，为每个核心提供完整的（必需的）data 数组以及计算整个results 数组的单个块并返回它们以在之后合并它们的任务。 （使用pool.apply_async(func, (data,))）

随着data 的大小进一步增加，我开始出现内存错误。根据我的分析，这是因为每个核心都有其单独的 data 数组，这意味着我的 RAM 中有 data 的原始 1 + # of cores 副本。为了减少这一点，我想我会尝试使用带有代理字典的管理器来处理data，然后每个核心都可以访问它。 （使用data_shared = manager.dict(data)）不幸的是，这非常慢，这大概就是不推荐它的原因。

我错过了一个明显的解决方案吗？我非常感谢任何想法。

Answer 1

这是一个使用整数列表的示例。

这段代码只是将一个包含 0、1、2、... 999 的 1000 个整数的数组分解为大小为 10 的块，最终得到 1000 / 10 = 100 个块，并从这些块中创建 100 个列表。这些块及其起始索引被提交给一个工作程序，该工作程序对列表的元素求和并返回起始索引和总和，然后用于计算总计。此处的起始索引未用于计算总计，但在另一种情况下它可能很有用。

from multiprocessing import Pool, Array
import itertools


def init_pool(arr):
    global data
    data = arr


def my_worker(tpl):
    index, results_chunk = tpl
    # sum items from results_chunk and data array:
    data_index = index
    the_sum = 0
    for item in results_chunk:
        the_sum += item + data[data_index]
        data_index += 1
    return index, the_sum


def get_chunks(arr, size):
    index = 0
    it = iter(arr)
    while True:
        x = tuple(itertools.islice(it, size))
        if not x:
            return
        yield index, list(x)
        index += size


# required for Windows and other platforms that do not have a fork() call:
if __name__ == '__main__':
    data = [x for x in range(1000, 2000)]
    data_sum = sum(data)
    arr = Array('i', data, lock=False)
    results = [x for x in range(1000)]
    results_sum = sum(results)
    print('Excpected total sum =', data_sum + results_sum)
    with Pool(initializer=init_pool, initargs=(arr,)) as pool:
        total = 0
        for index, the_sum in pool.imap_unordered(my_worker, get_chunks(results, 10), 3):
            total += the_sum
        print('Actual total sum =', total)

打印：

Excpected total sum = 1999000
Actual total sum = 1999000