numpy/scipy 的 Python 多处理速度慢

Question

我有一个处理器密集型任务，需要 13-20 小时才能完成，具体取决于机器。似乎是通过多处理库进行并行化的明显选择。问题是......我产生的进程越多，相同的代码就越慢。

每次迭代的时间（即运行 sparse.linalg.cg 所需的时间）：

183s 1 个进程

245s 2 个进程

312s 3 个进程

383s 4 个进程

当然，虽然 2 个进程在每次迭代中多花费 30% 多一点的时间，但它同时执行 2 个进程，所以它仍然稍微快一些。但我不希望实际的数学运算本身会变慢！这些计时器在多处理增加的任何开销之后才会启动。

这是我的代码的精简版。问题线是 sparse.linalg.cg 之一。 （我尝试过使用 MKL 与 OpenBLAS 之类的方法，并强制它们在单个线程中运行。还尝试手动生成进程而不是使用池。不走运。）

def do_the_thing_partial(iteration: int, iter_size: float, outQ : multiprocessing.Queue, L: int, D: int, qP: int, elec_ind: np.ndarray, Ic: int, ubi2: int,
                 K : csc_matrix, t: np.ndarray, dip_ind_t: np.ndarray, conds: np.ndarray, hx: float, dstr: np.ndarray):
    range_start = ceil(iteration * iter_size)
    range_end = ceil((iteration + 1) * iter_size)

    for rr in range(range_start, range_end):
        # do some things (like generate F from rr)
        Vfull=sparse.linalg.cg(K,F,tol=1e-11,maxiter=1200)[0] #Solve the system
        # do more things
        outQ.put((rr, Vfull))


def do_the_thing(L: int, D: int, qP: int, elec_ind: np.ndarray, Ic: int, ubi2: int,
                 K : csc_matrix, t: np.ndarray, dip_ind_t: np.ndarray, conds: np.ndarray, hx: float, dstr: np.ndarray):
    num_cores = cpu_count()
    iterations_per_process = (L-1) / num_cores  # 257 / 8 ?

    outQ = multiprocessing.Queue()

    pool = multiprocessing.Pool(processes=num_cores)

    [pool.apply_async(do_the_thing_partial,
                      args=(i, iterations_per_process, outQ, L, D, qP, elec_ind, Ic, ubi2, K, t, dip_ind_t, conds, hx, dstr),
                      callback=None)
     for i in range(num_cores)]

    pool.close()
    pool.join()

    for res in outQ:
        # combine results and return here

是我做错了什么，还是因为 sparse.linalg.cg 自身的优化而无法并行化？

谢谢！

Answer 1

这是一个如何使用Ray（用于并行和分布式 Python 的库）获得加速的示例。执行pip install ray（在Linux或MacOS上）后，您可以运行下面的代码。

在我的笔记本电脑上运行以下计算的串行版本（例如，执行 scipy.sparse.linalg.cg(K, F, tol=1e-11, maxiter=100) 20 次）需要 33 秒。为下面的代码计时以启动 20 个任务并获得结果需要 8.7 秒。我的笔记本电脑有 4 个物理内核，所以这几乎是 4 倍加速。

我对你的代码做了很多修改，但我认为我保留了它的精髓。

import numpy as np
import ray
import scipy.sparse
import scipy.sparse.linalg

# Consider passing in 'num_cpus=psutil.cpu_count(logical=True)'.
ray.init()

num_elements = 10**7
dim = 10**4

data = np.random.normal(size=num_elements)
row_indices = np.random.randint(0, dim, size=num_elements)
col_indices = np.random.randint(0, dim, size=num_elements)

K = scipy.sparse.csc_matrix((data, (row_indices, col_indices)))

@ray.remote
def solve_system(K, F):
    # Solve the system.
    return scipy.sparse.linalg.cg(K, F, tol=1e-11, maxiter=100)[0]

# Store the array in shared memory first. This is optional. That is, you could
# directly pass in K, however, this should speed it up because this way it only
# needs to serialize K once. On the other hand, if you use a different value of
# "K" for each call to "solve_system", then this doesn't help.
K_id = ray.put(K)

# Time the code below!

result_ids = []
for _ in range(20):
    F = np.random.normal(size=dim)
    result_ids.append(solve_system.remote(K_id, F))

# Run a bunch of tasks in parallel. Ray will schedule one per core.
results = ray.get(result_ids)

对ray.init() 的调用会启动 Ray 工作进程。对solve_system.remote 的调用将任务提交给工作人员。 Ray 默认情况下会为每个核心安排一个，但您可以通过@ray.remote(num_cpus=2) 指定特定任务需要更多资源（或更少资源）。您还可以指定 GPU 资源和其他自定义资源。

对solve_system.remote 的调用立即返回一个表示计算最终输出的ID，对ray.get 的调用获取ID 并检索计算的实际结果（因此ray.get 将等到任务完成执行）。

一些注意事项

• 在我的笔记本电脑上，scipy.sparse.linalg.cg 似乎将自己限制为单个内核，但如果不是，那么您应该考虑将每个工作进程固定到特定内核以避免工作进程之间的争用（您可以在 Linux 上执行此操作通过执行psutil.Process().cpu_affinity([i]) 其中i 是要绑定到的核心的索引。
• 如果所有任务都需要不同的时间，请确保您不只是在等待一项非常慢的任务。您可以通过从命令行运行 ray timeline 并在 chrome://tracing（在 Chrome 网络浏览器中）中可视化结果来检查这一点。
• Ray 使用shared memory object store 来避免每个worker 必须序列化和反序列化K 矩阵一次。这是一个重要的性能优化（尽管任务是否需要很长时间并不重要）。这主要有助于包含大型 numpy 数组的对象。它对任意 Python 对象没有帮助。这是通过使用Apache Arrow 数据布局启用的。你可以在this blog post阅读更多内容。

您可以在Ray documentation 中查看更多信息。请注意，我是 Ray 开发人员之一。