Question

我打算在Python中并行处理for循环，如下所示，用于处理大型数据数组。线程/核心/节点上的并行化如何适合此代码，以及如何实现它？任何建议表示赞赏。谢谢！

所有输入都是具有以下典型大小的NumPy数组：

vector_data (int64): 1M x 3
matrix (float64): 0.1M x 0.1M x 3

根据帖子的答案进行编辑：

对运行时性能的测试表明multiprocessing会导致运行速度显着下降以及对内存的更高要求。

from timeit import timeit
from multiprocessing import Pool

import numpy as np
from numba import jit

def OP():
    N = len(matrix_data)
    pop_array = np.zeros((N, N))
    for vector in vector_data:
        vector_2 = np.dot(vector, vector)
        pop_array += (np.exp(-vector_2) / vector_2
                      * np.cos(np.tensordot(matrix_data, vector, axes=([2], [0]))))
    return pop_array

def worker(vector):
    vector_2 = np.dot(vector, vector)
    return (np.exp(-vector_2) / vector_2
            * np.cos(np.tensordot(matrix_data, vector, axes=([2], [0]))))

def f1():
    N = len(matrix_data)
    pop_array = np.zeros((N, N))

    with Pool() as pool:
        results = pool.map(worker, vector_data)

    for res in results:
        pop_array += res

    return pop_array

def f2():
    N = len(matrix_data)
    pop_array = np.zeros((N, N))

    with Pool() as pool:
        for result in pool.imap(worker, vector_data):
            pop_array += result

    return pop_array

jit(parallel=True)
def f3():
    N = len(matrix_data)
    pop_array = np.zeros((N, N)) 
    for vector in vector_data:
        vector_2 = np.dot(vector, vector)
        pop_array += (np.exp(-vector_2) / vector_2
                      * np.cos(np.tensordot(matrix_data, vector, axes=([2], [0]))))
    return pop_array

max_vector_index = 150
vector_size = int(1E3)
matrix_size = int(1E2)

vector_shape = vector_size, 3
matrix_shape = matrix_size, matrix_size, 3

vector_data = np.random.randint(-max_vector_index, max_vector_index+1, vector_shape)
matrix_data = np.random.random(matrix_shape)

print(f'OP: {timeit(OP, number=10):.3e} sec')
print(f'f1: {timeit(f1, number=10):.3e} sec')
print(f'f2: {timeit(f2, number=10):.3e} sec')
print(f'f3: {timeit(f3, number=10):.3e} sec')

以下是样本运行的运行时成本：

vector_size = int(1E2)
matrix_size = int(1E1)

OP: 9.527e-02 sec
f1: 2.402e+00 sec (25.21x)
f2: 2.269e+00 sec (23.82x)
f3: 3.414e-02 sec (0.36x)

OP: 43.0 MiB
f1: 41.9 MiB (0.97x)
f2: 41.9 MiB (0.97x)

vector_size = int(1E3)
matrix_size = int(1E2)

OP: 1.420e+00 sec
f1: 1.448e+01 sec (10.20x)
f2: 2.051e+01 sec (14.44x)
f3: 1.213e+00 sec (0.86x)

OP: 43.4 MiB
f1: 119.0 MiB (2.74x)
f2: 43.8 MiB (1x)

vector_size = int(1E4)
matrix_size = int(1E3)

OP: 5.116e+02 sec
f1: 8.902e+02 sec (1.74x)
f2: 6.509e+02 sec (1.27x)

OP: 73.9 MiB
f1: 76402.1 MiB (1033x)
f2: 209.7 MiB (2.84x)

Answer 1

您可以使用multiprocessing Pool。然后，您可以使用map方法在可迭代对象上运行函数。因此，您可以首先创建要传递给工作人员的函数，以对可迭代对象的每个元素进行处理：

def worker(vector):
    vector_2 = np.dot(vector, vector)
    return (np.exp(-vector_2) / vector_2
            * np.cos(np.tensordot(matrix, vector, axes=([2], [0]))))

现在，您可以创建Pool在每个矢量上运行此功能。它将返回结果列表，然后我们可以将这些结果添加到pop_array中。像这样：

from multiprocessing import Pool

def par_fun(vector_data, matrix):
    N = len(matrixA)
    pop_array = np.zeros((N, N))

    with Pool() as pool:
        results = pool.map(worker, vector_data)

    for res in results:
        pop_array += res

    return pop_array

另一种更整洁的方式是使用imap。从文档中：

请注意，这可能会导致很长的可迭代项占用大量内存。考虑将imap()或imap_unordered()与明确的 chunksize 一起使用选择以提高效率。

也：

chunksize 参数与map()使用的参数相同方法。对于很长的可迭代对象，使用较大的 chunksize 值可以与使用默认值1相比，使工作完成的速度快得多。

因此您可以使用以下代码：

def par_fun(vector_data, matrix):
    N = len(matrixA)
    pop_array = np.zeros((N, N))

    pool_size = None
    chunksize = 1

    with Pool(pool_size) as pool:
        for result in pool.imap(worker, vector_data, chunksize=chunksize):
            pop_array += result

    return pop_array

并使用不同的pool_size和chunksize值来达到最佳效果。

另一个选择是使用线程而不是进程。流程具有创建和维护的开销，这可能会影响运行时。要将代码更改为使用线程，只需将导入更改为使用dummy包装器：

from multiprocessing.dummy import Pool

其余代码保持不变

Python中线程/核心/节点上的并行for循环

1 个答案: