Python多进程未达到预期加速效果

Question

我正在尝试使用Python的multiprocessing.Pool模块来优化我的代码，但我得到的速度提升并没有我预期的那么明显。

我主要的工作是计算大量向量和一个固定的大稀疏矩阵之间的矩阵-向量乘法。下面是一个简单的例子，它能完成我需要的功能，但使用的是随机矩阵。

import time
import numpy as np
import scipy.sparse as sp

def calculate(vector, matrix = None):
    for i in range(50):
        v = matrix.dot(vector)
    return v

if __name__ == '__main__':
    N = 1e6
    matrix = sp.rand(N, N, density = 1e-5, format = 'csr')
    t = time.time()
    res = []
    for i in range(10):
        res.append(calculate(np.random.rand(N), matrix = matrix))    
    print time.time() - t

这个方法大约需要30秒才能完成。

现在，由于results中每个元素的计算都不依赖于其他计算的结果，所以自然会想到并行计算会加快这个过程。我的想法是创建4个进程，如果每个进程都做一些计算，那么所有进程完成的时间应该会减少大约4倍。为此，我写了以下代码：

import time
import numpy as np
import scipy.sparse as sp
from multiprocessing import Pool
from functools import partial

def calculate(vector, matrix = None):
    for i in range(50):
        v = matrix.dot(vector)
    return v

if __name__ == '__main__':
    N = 1e6
    matrix = sp.rand(N, N, density = 1e-5, format = 'csr')

    t = time.time()
    input = []
    for i in range(10):
        input.append(np.random.rand(N))
    mp = partial(calculate, matrix = matrix)
    p = Pool(4)
    res = p.map(mp, input)
    print time.time() - t

我的问题是，这段代码运行的时间稍微超过了20秒，所以我甚至没有提高2倍的性能！更糟糕的是，即使池中有8个进程，性能也没有改善！你知道为什么没有实现速度提升吗？

---

注意：我实际的方法需要更长的时间，而且输入向量存储在一个文件中。如果我把文件分成4个部分，然后手动为每个文件分别运行我的脚本，每个进程完成的时间是处理整个文件的四分之一（这是预期的）。我很困惑，为什么这种明显可行的速度提升在使用multiprocessing.Pool时没有发生。

---

补充：我刚刚发现了一个相关的问题，链接是Multiprocessing.Pool使Numpy矩阵乘法变慢。不过我还需要检查一下。

Answer 1

试试这个：

p = Pool(4)
for i in range(10):
    input = np.random.rand(N)
    p.apply_async(calculate, args=(input, matrix)) # perform function calculate as new process with arguments input and matrix

p.close()  
p.join() # wait for all processes to complete

我怀疑“partial”对象和map可能导致了阻塞的情况。（不过我从来没用过partial，所以不太了解。）

“apply_async”（或者“map_async”）是多进程的方法，它们特别设计成不会阻塞。（可以参考这个链接：Python multiprocessing.Pool: 什么时候使用apply、apply_async或map？）

一般来说，对于像这样的“非常适合并行处理的问题”，apply_async对我来说效果很好。

编辑：

我通常在完成后会把结果写入MySQL数据库——我提供的实现如果不是这种方式就不适用了。如果你想根据列表中的顺序来跟踪每个条目，那么“map”可能是正确的选择，但我对“partial”对象还是有点怀疑。