如何从多处理池映射返回 numpy 数组？

Question

我有一个函数“func”，它接收一个包含名为“matrix_image”的图像行数的列表：

list_rows = range(N_rows)

然后，计算是在 func 内部完成的，我得到了代表图像的结果矩阵的新行。

def func(list_rows):
    new_row = numpy.empty(N_columns)
    ....
    ....# some computation
    ....
    return new_row


matrix_image = pool.map(func, list_n_rows )

“new_row”在func正确计算，因为我在调试器中看到值，但结果 matrix_image 具有形状：(list_n_rows,)。 填充了None 的所有值。

如何从func返回（1D numpy array）行以附加到结果矩阵（2D numpy array）中？

Answer 1

您可以使用 multiprocessing 的 RawArray 功能，在该功能中，您定义需要从进程访问的变量，然后在启动进程之前将其作为 RawArray 进行访问，然后在进程完成后将其作为重新整形的 numpy 数组进行访问。

下面是一个例子：

import numpy as np
import multiprocessing as mp

n_elements = 1000 # how many elements your numpy should have

def myProc( shared_var ):
    '''
    here you convert your shared variable from mp.RawArray to numpy
    then treat it as it is numpy array e.g. fill it in with some 
    random numbers for demonstration purpose
    '''
    var = np.reshape( np.frombuffer( shared_var, dtype=np.uint32 ), -1 )
    for i in range( n_elements ):    
        var[i] = np.random.randint( 0, 2**16, 1 )
    print( 'myProc var.mean() = ', var.mean() )               
                                  
#buffer that contains the memory
mp_var = mp.RawArray( 'i', n_elements )
p = mp.Process( target=myProc, args=(mp_var,) )
p.start()                                      
p.join()

#after the process has ended, you convert the buffer that was passed to it
var = np.reshape( np.frombuffer( mp_var, dtype=np.uint32 ), -1)
#and again, you can treat it like a numpy array
print( '   out var.mean() = ',var.mean() )                                            

输出是：

myProc var.mean() =  32612.403
       var.mean() =  32612.403

希望有帮助！

请注意，如果您从并发进程访问此缓冲区，则需要组织适当的锁定机制，以免两个进程同时修改同一块内存。