在多进程中使用共享内存的numpy数组

Question

我想在共享内存中使用一个numpy数组，以便和多进程模块一起使用。问题在于，我希望它能像一个numpy数组那样使用，而不仅仅是一个ctypes数组。

from multiprocessing import Process, Array
import scipy

def f(a):
    a[0] = -a[0]

if __name__ == '__main__':
    # Create the array
    N = int(10)
    unshared_arr = scipy.rand(N)
    arr = Array('d', unshared_arr)
    print "Originally, the first two elements of arr = %s"%(arr[:2])

    # Create, start, and finish the child processes
    p = Process(target=f, args=(arr,))
    p.start()
    p.join()

    # Printing out the changed values
    print "Now, the first two elements of arr = %s"%arr[:2]

这样会产生类似以下的输出：

Originally, the first two elements of arr = [0.3518653236697369, 0.517794725524976]
Now, the first two elements of arr = [-0.3518653236697369, 0.517794725524976]

这个数组可以用ctypes的方式访问，比如说 arr[i] 是可以理解的。但是，它并不是一个真正的numpy数组，所以我不能进行像 -1*arr 或者 arr.sum() 这样的操作。我想解决这个问题的方法可能是把ctypes数组转换成numpy数组。不过（除了我不知道怎么做），我觉得这样就不再是共享的了。

看起来应该有一个标准的解决方案来处理这个常见的问题。

Answer 1

Array对象有一个叫做get_obj()的方法，这个方法可以返回一个ctypes数组，它提供了一种缓冲区接口。我觉得下面的代码应该可以正常运行...

from multiprocessing import Process, Array
import scipy
import numpy

def f(a):
    a[0] = -a[0]

if __name__ == '__main__':
    # Create the array
    N = int(10)
    unshared_arr = scipy.rand(N)
    a = Array('d', unshared_arr)
    print "Originally, the first two elements of arr = %s"%(a[:2])

    # Create, start, and finish the child process
    p = Process(target=f, args=(a,))
    p.start()
    p.join()

    # Print out the changed values
    print "Now, the first two elements of arr = %s"%a[:2]

    b = numpy.frombuffer(a.get_obj())

    b[0] = 10.0
    print a[0]

当运行这段代码时，它会打印出a的第一个元素，现在是10.0，这表明a和b其实是同一块内存的两个不同视图。

为了确保在多处理器环境下也能安全使用，我认为你需要使用Array对象a上的acquire和release方法，以及它内置的锁，确保所有的访问都是安全的（不过我对多处理器模块不是很专业）。

Answer 2

虽然之前的回答都不错，但其实这个问题有一个更简单的解决办法，只要满足两个条件：

1. 你使用的是一个符合POSIX标准的操作系统（比如Linux或Mac OSX）；
2. 你的子进程只需要只读访问共享的数组。

在这种情况下，你不需要特意去设置变量为共享，因为子进程是通过“分叉”（fork）创建的。分叉出来的子进程会自动共享父进程的内存空间。在Python的多进程处理中，这意味着它会共享所有的模块级变量；不过要注意，这对你明确传递给子进程的参数或者在multiprocessing.Pool中调用的函数是不适用的。

下面是一个简单的例子：

import multiprocessing
import numpy as np

# will hold the (implicitly mem-shared) data
data_array = None

# child worker function
def job_handler(num):
    # built-in id() returns unique memory ID of a variable
    return id(data_array), np.sum(data_array)

def launch_jobs(data, num_jobs=5, num_worker=4):
    global data_array
    data_array = data

    pool = multiprocessing.Pool(num_worker)
    return pool.map(job_handler, range(num_jobs))

# create some random data and execute the child jobs
mem_ids, sumvals = zip(*launch_jobs(np.random.rand(10)))

# this will print 'True' on POSIX OS, since the data was shared
print(np.all(np.asarray(mem_ids) == id(data_array)))

Answer 3

为了补充@unutbu（现在无法访问）和@Henry Gomersall的回答，你可以使用 shared_arr.get_lock() 来在需要的时候同步访问：

shared_arr = mp.Array(ctypes.c_double, N)
# ...
def f(i): # could be anything numpy accepts as an index such another numpy array
    with shared_arr.get_lock(): # synchronize access
        arr = np.frombuffer(shared_arr.get_obj()) # no data copying
        arr[i] = -arr[i]

示例

import ctypes
import logging
import multiprocessing as mp

from contextlib import closing

import numpy as np

info = mp.get_logger().info

def main():
    logger = mp.log_to_stderr()
    logger.setLevel(logging.INFO)

    # create shared array
    N, M = 100, 11
    shared_arr = mp.Array(ctypes.c_double, N)
    arr = tonumpyarray(shared_arr)

    # fill with random values
    arr[:] = np.random.uniform(size=N)
    arr_orig = arr.copy()

    # write to arr from different processes
    with closing(mp.Pool(initializer=init, initargs=(shared_arr,))) as p:
        # many processes access the same slice
        stop_f = N // 10
        p.map_async(f, [slice(stop_f)]*M)

        # many processes access different slices of the same array
        assert M % 2 # odd
        step = N // 10
        p.map_async(g, [slice(i, i + step) for i in range(stop_f, N, step)])
    p.join()
    assert np.allclose(((-1)**M)*tonumpyarray(shared_arr), arr_orig)

def init(shared_arr_):
    global shared_arr
    shared_arr = shared_arr_ # must be inherited, not passed as an argument

def tonumpyarray(mp_arr):
    return np.frombuffer(mp_arr.get_obj())

def f(i):
    """synchronized."""
    with shared_arr.get_lock(): # synchronize access
        g(i)

def g(i):
    """no synchronization."""
    info("start %s" % (i,))
    arr = tonumpyarray(shared_arr)
    arr[i] = -1 * arr[i]
    info("end   %s" % (i,))

if __name__ == '__main__':
    mp.freeze_support()
    main()

如果你不需要同步访问，或者你自己创建了锁，那么 mp.Array() 就没必要了。在这种情况下，你可以使用 mp.sharedctypes.RawArray。