限制numpy中的线程数

Question

看来我的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个线程进行矩阵乘法.

Answer 1

尝试设置以下所有内容:

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

有时候看到引入多线程的确切位置有点棘手.

Answer 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这一点。他们提交更改后，我将对其进行更新。

Answer 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评论转发的答案.)

Answer 4

After trying a number of the solutions above without luck, I found a reference to threadpoolctl in the Numpy docs. This worked and it can be used even if numpy is already imported.

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

Just make sure to use the user_api which is listed when you do:

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