时间戳

首页 > 解决方案 > numpy matmul 是否并行化以及如何停止它？

python - numpy matmul 是否并行化以及如何停止它？

问题描述

在我的脚本执行期间查看资源监视器，我注意到我的 PC 的所有内核都在工作，即使我没有实现任何形式的多处理。试图查明原因，我发现在使用numpy's时代码是并行化的matmult（或者，如下例所示，二元运算符@）。

import numpy as np

A = np.random.rand(10,500)
B = np.random.rand(500,50000)
while True:
    _ = A @ B

看看这个问题，原因似乎是numpy调用BLAS/LAPACK了确实并行化的例程。

尽管我的代码运行得更快并使用了所有可用资源，这很好，但当我在PBS队列管理器管理的共享集群上提交代码时，这给我带来了麻烦。与集群 IT 经理一起，我们注意到即使我在集群节点上请求 N 个 CPU，numpy仍然会产生与节点上 CPU 数量相等的线程数。

这导致节点过载，因为我使用的 CPU 比分配给我的 CPU 多。

有没有办法“控制”这种行为并告诉numpy它可以使用多少个 CPU？

标签： pythonnumpymultiprocessing

解决方案

您可以尝试使用threadpoolctl. 有关详细信息，请参阅自述文件。不过，在使用之前，我建议您查看“已知限制”部分。

该自述文件中的引用

Python 帮助程序，用于限制用于科学计算和数据科学（例如 BLAS 和 OpenMP）的常用本机库的线程池支持中使用的线程数。

对底层线程池大小的精细控制在涉及嵌套并行的工作负载中非常有用，可以缓解超额订阅问题。

该自述文件中的代码片段

from threadpoolctl import threadpool_limits
import numpy as np


with threadpool_limits(limits=1, user_api='blas'):
    # In this block, calls to blas implementation (like openblas or MKL)
    # will be limited to use only one thread. They can thus be used jointly
    # with thread-parallelism.
    a = np.random.randn(1000, 1000)
    a_squared = a @ a

推荐阅读