Python 中多處理返回值的快速整形

Question

我有一個 function 和列表值返回值，我在 Python 中進行多重處理，我需要在最后將它們連接到一維列表。 以下是演示示例代碼：

import numpy as np
import multiprocessing as mp
import random as rd

N = 4
L = list(range(0, N))

def F(x):
    a = []
    b = []
    for t in range(0,2):
        a.append('a'+str(t*x))
        b.append('b'+str(t*x))
    return a, b

pool = mp.Pool(mp.cpu_count())
a,b = zip(*pool.map(F, L))
pool.close()

print(a)
print(b)

A = np.concatenate(a)
B = np.concatenate(b)
              
print(A)
print(B)

用於說明的 output 是：

(['a0', 'a0'], ['a0', 'a1'], ['a0', 'a2'], ['a0', 'a3'])
(['b0', 'b0'], ['b0', 'b1'], ['b0', 'b2'], ['b0', 'b3'])
['a0' 'a0' 'a0' 'a1' 'a0' 'a2' 'a0' 'a3']
['b0' 'b0' 'b0' 'b1' 'b0' 'b2' 'b0' 'b3']

問題是我正在處理的列表L非常大，並且最后的連接需要大量時間，這大大減少了串行處理的優勢。

是否有一些巧妙的方法可以避免串聯，或者有更快的方法來執行串聯？ 我一直在擺弄隊列，但這似乎有點慢。

注意：這似乎是一個與Add result from multiprocessing into array類似的問題。

Answer 1

如果所需的 output 是適合創建scipy.sparse.coo_matrix的輸入，我會采取非常不同的方法：不返回任何內容，只需創建可以直接修改的共享對象。

創建coo_matrix需要的是數據值數組、數據行數組和數據列數組（除非您已經有另一個稀疏矩陣/密集矩陣）。 我將創建 3 個共享的 arrays ，每個進程都可以使用L中的每個條目的索引將結果直接轉儲到其中。 這甚至允許亂序執行，因此您可以使用imap_unordered來提高速度：

from multiprocessing.pool import Pool
from multiprocessing.sharedctypes import RawArray
from random import random, randint # bogus data for testing
import numpy as np

from ctypes import c_int, c_float
from scipy.sparse import coo_matrix

#pool worker globals are only global to that process
worker_globals = {}

def init_worker(data_array, row_array, col_array):
    worker_globals['data'] = np.frombuffer(data_array, dtype=c_float)
    worker_globals['row'] = np.frombuffer(row_array, dtype=c_int)
    worker_globals['col'] = np.frombuffer(col_array, dtype=c_int)

def worker_func(tup):
    i, x = tup #enumerate returns a tuple with the index then the value
    #don't bother with mutexes because we only ever write to array[i] once from a single process
    worker_globals['data'][i] = random() #calculate your data, row, and column, and write back to the shared arrays
    worker_globals['row'][i] = x
    worker_globals['col'][i] = randint(0,1000)

if __name__ == "__main__":
    L = list(range(100, 0, -1)) #some data in L
    
    data_array = RawArray(c_float, len(L))
    row_array = RawArray(c_int, len(L))
    col_array = RawArray(c_int, len(L))
    
    with Pool(initializer=init_worker, initargs=(data_array, row_array, col_array)) as p:
        for _ in p.imap_unordered(worker_func, enumerate(L)):
            pass
    d = np.frombuffer(data_array, dtype=c_float)
    r = np.frombuffer(row_array, dtype=c_int)
    c = np.frombuffer(col_array, dtype=c_int)
    mat = coo_matrix((d, (r, c)), shape=(1000, 1000))

順便說一句：在使用多處理時，您也絕對應該始終使用if __name__ == "__main__": 它在任何地方都被建議，並且在 windows 上需要。