python multiprocessing.Pool类中的全局变量

Question

我有一个相对简单的可并行化问题，导致太多问题无法有效实施。 我程序的核心是两个矩阵和两个向量，对于我要执行的两个计算中的每一个，一个矩阵和一个向量。

在代码中，这意味着我有

import numpy as np
matrices = dict([("type1", np.random.rand(10,10)), ("type2", np.random.rand(10,10))])
vectors = dict([("type1", np.random.rand(10)), ("type2", np.random.rand(10))])

我想做的（不是真的，但在简化的情况下）是这样的：

对于每种类型，我都有大量的向量：

input_vectors = [np.random.rand(10) for i in range(1000)]

我想计算A*v+b ，其中A是矩阵， b是每种类型的向量。

因此，执行我需要的单线程代码是

def f(input_vector, matricx, vector):
    return np.dot(matrix, input_vector) + vector

results = {}
for type in ['type1', 'type2']:
    results[type] = []
    for input_vector in input_vectors:
        results.append(f(input_vector, matrices[type], vectors[type]))

但是，我想并行执行此操作。 但是，我不知道如何解决我要映射到向量列表上的函数将其输入当作不仅仅是向量的问题。

我想写点东西

from multiprocessing import Pool
p = Pool(4)
for type in types:
    p.map(lambda x: f(x, matrices[type], vectors[type] , input_vectors))

但是，这不起作用，因为无法腌制lambda函数。 起作用的一件事是将要与之相乘的矩阵附加到每个向量，但是从内存角度来看，这当然是不可行的。

关于如何优雅地解决我的难题的任何想法？

---

我想要的是池中的每个元素都有一个矩阵和向量的副本，它必须与之相乘，但是我不知道如何在multiprocessing做到这一点。

Answer 1

使用functools.partial将多个参数传递给map ：

def f(matrix, vector, input_vector):
    return np.dot(matrix, input_vector) + vector

results = {}
for type_ in types:
    func = partial(f, matrices[type_], vectors[type_])
    results[type_] = p.map(func, input_vectors)

如果您希望在启动Pool时将整个matrices和vectors列表传递给每个孩子，然后在调用map时仅传递type ，您也可以这样做。 使用initializer / initargs参数进行multiprocessing.Pool initargs传递列表，然后将它们全局设置为initializer函数内部。 这将使它们在每个子进程中都具有全局性：

matrices = vectors = None

def init(_matrices, _vectors):
    global matrices, vectors
    matrices = _matrices
    vectors = _vectors


def f(type_, input_vector):
    return np.dot(matrices[type_], input_vector) + vectors[type_]

def main():
    # <declare matrices, vectors, input_vectors here>
    p = multiprocessing.Pool(initializer=init, 
                             initargs=(matrices, vectors))
    results = {}
    for type_ in ['type1', 'type2']:
        func = partial(f, type_)
        results[type_] = p.map(func, input_vectors)

if __name__ == "__main__":
    main()