限制numpy中的线程数

问题描述：

看来我的 numpy 库正在使用 4 个线程，并且设置OMP_NUM_THREADS=1并不能停止这种情况。

numpy.show_config()给出了以下结果：

atlas_threads_info:
    libraries = ['lapack', 'ptf77blas', 'ptcblas', 'atlas']
    library_dirs = ['/usr/lib64/atlas']
    define_macros = [('ATLAS_INFO', '"\\\"3.8.4\\\""')]
    language = f77
    include_dirs = ['/usr/include']
blas_opt_info:
    libraries = ['ptf77blas', 'ptcblas', 'atlas']
    library_dirs = ['/usr/lib64/atlas']
    define_macros = [('ATLAS_INFO', '"\\\"3.8.4\\\""')]
    language = c
    include_dirs = ['/usr/include']
atlas_blas_threads_info:
    libraries = ['ptf77blas', 'ptcblas', 'atlas']
    library_dirs = ['/usr/lib64/atlas']
    define_macros = [('ATLAS_INFO', '"\\\"3.8.4\\\""')]
    language = c
    include_dirs = ['/usr/include']
openblas_info:
  NOT AVAILABLE
lapack_opt_info:
    libraries = ['lapack', 'ptf77blas', 'ptcblas', 'atlas']
    library_dirs = ['/usr/lib64/atlas']
    define_macros = [('ATLAS_INFO', '"\\\"3.8.4\\\""')]
    language = f77
    include_dirs = ['/usr/include']

所以我知道它正在使用 blas，但我不知道如何让它使用 1 个线程进行矩阵乘法。

解决方案 1：

有一些常见的多 CPU 库可用于数值计算，包括 NumPy 内部。在运行脚本之前，您可以设置一些环境标志来限制它们使用的 CPU 数量。

尝试设置以下所有内容：

export MKL_NUM_THREADS=1
export NUMEXPR_NUM_THREADS=1
export OMP_NUM_THREADS=1

有时，要看清楚多线程究竟在哪里引入是有点困难的。

其他答案显示了其他库的环境标志。它们也可能有效。

解决方案 2：

环境变量不止上述 3 个。以下是环境变量的完整列表以及使用该变量来控制其生成的线程数的包。请注意，您需要在执行以下操作之前设置这些变量import numpy：

OMP_NUM_THREADS: openmp,
OPENBLAS_NUM_THREADS: openblas,
MKL_NUM_THREADS: mkl,
VECLIB_MAXIMUM_THREADS: accelerate,
NUMEXPR_NUM_THREADS: numexpr

因此在实践中你可以这样做：

import os
os.environ["OMP_NUM_THREADS"] = "4" # export OMP_NUM_THREADS=4
os.environ["OPENBLAS_NUM_THREADS"] = "4" # export OPENBLAS_NUM_THREADS=4 
os.environ["MKL_NUM_THREADS"] = "6" # export MKL_NUM_THREADS=6
os.environ["VECLIB_MAXIMUM_THREADS"] = "4" # export VECLIB_MAXIMUM_THREADS=4
os.environ["NUMEXPR_NUM_THREADS"] = "6" # export NUMEXPR_NUM_THREADS=6

请注意，自 2018 年 11 月起，Numpy 开发人员正在努力实现这一点，在您这样做之后import numpy也是如此。一旦他们提交这些更改，我就会更新这篇文章。

解决方案 3：

关于从 python 脚本中而不是在 bash 提示符下执行此操作，按照此线程， 您可以执行以下操作（与上面的答案相同的命令）：

import os
os.environ["MKL_NUM_THREADS"] = "1" 
os.environ["NUMEXPR_NUM_THREADS"] = "1" 
os.environ["OMP_NUM_THREADS"] = "1" 

但你必须在执行之前import numpy把它放好。显然，numpy 只在导入时检查这一点。

（这是根据上面@kηives 的评论重新发布的答案。）

解决方案 4：

在尝试了上述几种解决方案却无济于事之后，我在 Numpy文档中找到了对threadpoolctl的引用。这很有效，即使已经导入了 numpy 也可以使用它。

from threadpoolctl import threadpool_limits

with threadpool_limits(limits=1, user_api='blas'):
  # single threaded numpy code...

只要确保使用user_api列出的内容即可：

from threadpoolctl import threadpool_info
from pprint import pprint
import numpy
pprint(threadpool_info())

解决方案 5：

我能够通过以下方式在运行时修复此问题：

import mkl
mkl.set_num_threads(1)

我使用以下代码来使此代码片段不太可能在脚本/包中引起问题：

try:
    import mkl
    mkl.set_num_threads(1)
except:
    pass