Python 多处理抛出 Killed : 9

Question

我正在尝试使用 multiprocessing 来加速我将 2000 个形状数组 (76, 76) 平铺到 3D 数组并应用比例因子的函数。

当图 block 数量少于 200 个时它工作正常，但当它大于 200 个时我得到一个 Killed: 9 并且我需要能够处理 1000 个图 block 的订单。

这是代码的简化版本:

from functools import partial
from multiprocessing.pool import ThreadPool
from multiprocessing import cpu_count
import numpy as np

def func_A(data, scale, N):
    """Tile the data N times and scale it"""
    arr = np.tile(data, (N, 1, 1))
    arr *= scale
    return arr

def func_B(N=4):
    """Create scaled arrays"""
    # Make data
    data = np.random.normal(size=(2000, 76, 76))

    # Make scales
    scales = np.arange(2000)

    # Multiprocess into tiled arrays
    pool = ThreadPool(cpu_count())
    func = partial(func_A, N=N)
    inpt = list(zip(data, scales))
    results = np.asarray(pool.starmap(func, inpt), dtype=np.float64)
    pool.close()
    pool.join()

    return results.swapaxes(0, 1)

所以它对 func_B(4) 没问题，但对 func_B(500) 就死了。

我知道我正在用如此大的数组对 Python 的内存征税，但是让 func_B 与大的 N 一起工作的最佳方法是什么...最好是快速？我使用 multiprocessing 错了吗？我是否应该完全使用其他东西，例如Dask、Numba、Cython 等？

如有任何帮助，我们将不胜感激。谢谢!

Answer 1

我不完全确定你计算的目的是什么，但下面的内容似乎在 dask 中完成了这项工作

import dask.array as da
import numpy as np

# Make data
data = da.random.normal(size=(2000, 76, 76), chunks=(2000, 76, 76))

# Make scales
scales = np.arange(2000)
N = 500
out = da.repeat(data, N, axis=0).reshape((N, 2000, 76, 76)) * scales.reshape((1, 2000, 1, 1))
out = out.sum(axis=0).compute()

保持工作内存 <~5GB 并使用大部分内核。