python multiprocessing.Pool类中的全局变量

Question

I have a relatively simply parallelizable problem that is causing way too much problems for me to implement effectively. What is at the core of my program are two matrices and two vectors, one matrix and vector for each of my two calculations I wish to perform.

In code, that means I have

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))])

What I want to do (not really, but in the simplified case) is this:

I have a very large list of vectors for each type:

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

and I want to calculate A*v+b where A is the matrix and b is the vector for each type.

The single-thread code that does what I need is therefore

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]))

However, I want to do this in parallel. However, I do not know how to solve the problem that the function that I want to map over the list of vectors takes as its input more than just the vectors.

I want to write something like

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

However, that does not work because the lambda function cannot be pickled. One thing that does work is to append the matrix I want to multiply with to each vector, but that of course is not feasible memory-wise.

Any thoughts on how to elegantly solve my conundrum?

---

What I would like is for each element of the pool to have a copy of the matrix and vector it has to multiply with, but I do not know how to do this in multiprocessing .

Answer 1

Use functools.partial to pass multiple arguments to 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)

If you prefer to pass the entire matrices and vectors lists to each child when you start up the Pool , and then just pass the type when you call map , you can do that, too. Use the initializer / initargs arguments to multiprocessing.Pool to pass the lists, and then make them global inside the initializer function. That will make them global inside each child process:

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()