多处理多个大型 numpy 数组作为共享内存

Question

我有多个大型 numpy 数组：

x1=np.zeros((4000,4000))
x2=np.zeros((4000,4000))
x3=np.zeros((4000,4000))
.
.
.
xn=np.zeros((4000,4000))

我想用这些数组并行执行一个函数。 因为每个数组都独立于其他数组，所以我想我可以使用 shared_memory 这样子进程就不会腌制数据。

是否可以创建一个包含 3 个大 numpy 数组的大“共享变量”？

在子进程内部，我想直接写入这些数组（不腌制它们）。

我认为 a 会将 idx (0,1,2...n) 参数传递给子进程以引用 x1,x2,x3...xn 数组？

这可能吗？我认为一个数组不是问题，但是多个数组多处理让我有点困惑。

谢谢。

Answer 1

这就是使用共享内存数组的方法。

import numpy as np
import ctypes
from multiprocessing.sharedctypes import RawArray
from multiprocessing.pool import Pool

def main():
    n = ...  # Number of arrays
    # Make arrays
    x1 = np.zeros((4000, 4000), dtype=np.float64)
    x2 = np.zeros((4000, 4000), dtype=np.float64)
    # ...
    xn = np.zeros((4000, 4000), dtype=np.float64)
    # Make big array of shared memory (ctype must match array type)
    array_mem = RawArray(ctypes.c_double, n * 4000 * 4000)
    arr = np.frombuffer(array_mem, dtype=np.float64).reshape(n, 4000, 4000)
    arr[0] = x1
    arr[1] = x2
    # ...
    arr[n - 1] = xn
    # Process array in a pool of processes
    with Pool(initializer=init_process, initargs=(array_mem, arr.shape)) as p:
        p.map(process_array, range(n))
    # The array has been processed
    # ...
    print(*arr[:, :2, :3], sep='\n')
    # [[0. 0. 0.]
    #  [0. 0. 0.]]
    # [[100. 100. 100.]
    #  [100. 100. 100.]]
    # [[200. 200. 200.]
    #  [200. 200. 200.]]
    # ...

# Global values for subprocesses
process_array_mem = None
process_array_shape = None

# Process initializer saves memory pointer and array shape
def init_process(array_mem, array_shape):
    global process_array_mem, process_array_shape
    process_array_mem = array_mem
    process_array_shape = array_shape

def process_array(array_idx):
    # Create array from shared memory
    arr = np.frombuffer(process_array_mem, dtype=np.float64).reshape(process_array_shape)
    # Pick array for this process
    process_array = arr[array_idx]
    # Do processing
    process_array += 100 * array_idx

if __name__ == '__main__':
    main()

在上面的代码中，我使用n = ... 将数组的数量设置为在您的情况下具有的任何值，但是如果您将其更改为n = 3 并将sn-p 保存为文件，您可以运行它并看到结果。初始化器和全局值部分可能有点混乱，但问题是array_mem 必须由子进程继承，这意味着我不能将它作为另一个参数与map 一起传递，我认为这是最简单的使用方式它。