如何在 python 中并行化以下代码片段？

Question

我有一堆仅按行执行的矩阵乘法运算。我想知道如何通过并行化加速计算：

data = np.random.randint(1, 100, (100000, 800))
indices_1 = np.equal(data, 1)
A = np.zeros((100000, 100))
B = np.random.randn(800, 100)

for i in range(100000):
   ones = indices_1[i]
   not_ones = ~indices_1[i]
   B_ones = B[ones]
   B_not_ones = B[not_ones]
   A[i] = (data[i][not_ones] @ B_not_ones) @ np.linalg.inv(B_not_ones.T @ B_not_ones)  
   data[i][ones] = A[i] @ B_ones.T
    

我尝试了多处理器，但由于某种原因，它的性能并不比顺序处理器好。这是我的多处理器实现：

from multiprocessing.pool import ThreadPool, Pool
pool = ThreadPool() # can also use Pool

def f(i):
   ones = indices_1[i]
   not_ones = ~indices_1[i]
   B_ones = B[ones]
   B_not_ones = B[not_ones]
   A[i] = (data[i][not_ones] @ B_not_ones) @ np.linalg.inv(B_not_ones.T @ B_not_ones)  
   data[i][ones] = A[i] @ B_ones.T


pool.map(f, range(100000))

两者的运行时间相同（大约 32 秒）。 concurrent.futures 等其他并行化方法并没有提高运行时间（如下所示）：

with concurrent.futures.ThreadPoolExecutor() as executor:
     result = executor.map(f, range(100000))

我也尝试申请dask，但无法让他们的框架在我的情况下工作。任何帮助都感激不尽！谢谢！

Answer 1

import numpy as np
import multiprocessing as mp


data = list(np.random.randint(1, 100, (100000, 800)))
indices_1 = np.equal(data, 1)
A = list(np.zeros((100000, 100)))
B = np.random.randn(800, 100)


def f(data, A, i):
    ones = indices_1[i]
    not_ones = ~indices_1[i]
    B_ones = B[ones]
    B_not_ones = B[not_ones]
    A[i] = (data[i][not_ones] @ B_not_ones) @ np.linalg.inv(B_not_ones.T @ B_not_ones)
    data[i][ones] = A[i] @ B_ones.T

with mp.Manager() as manager:
    data_global = manager.list(data)
    A_global = manager.list(A)

    with mp.Pool() as p:
        results = [ p.apply_async(f, (data_global, A_global, i,)) for i in range(100000) ]
        for i in results:
            i.wait()

    data_global = list(data_global)
    A_global = list(A_global)