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

Question

在执行脚本期间查看资源监视器，我注意到我的 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？

Answer 1

您可以尝试使用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