logo标签列表

限制numpy中线程的数量

pythonmultithreadingnumpy
61

看起来我的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，但是我不知道如何让它只使用一个线程进行矩阵乘法。

3
Atlas在编译时定义线程数量。 - jfs
5个回答
75

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

试着设置以下所有内容：

export MKL_NUM_THREADS=1
export NUMEXPR_NUM_THREADS=1
export OMP_NUM_THREADS=1
有时候很难看出多线程是在哪里引入的。
其他答案展示了其他库的环境标志。它们也可能有效。
11
需要更多详细信息。 - Mohamad Elmasri
1
在Mac上，export VECLIB_MAXIMUM_THREADS=1; 参考performance-of-numpy-with-different-blas-implementations。注意：man Accelerate说：“VECLIB_MAXIMUM_THREADS的值可能会被库缓存并重复使用；如果您需要确保单线程执行，则应在调用任何Accelerate之前设置VECLIB_MAXIMUM_THREADS”。 - denis
1
从脚本内部来看，最好的方法是什么？我是否只需运行 os.system("export OMP_NUM_THREADS=1")？我已经阅读了关于它是否能够正确工作的混合评论，即使它确实可以，那是否是最好的方法呢？出于某种原因，这似乎有点不正规。 - seth127
34
如果还有其他人和我一样有同样的问题，可以参考这个帖子 http://numpy-discussion.10968.n7.nabble.com/Set-threads-from-within-python-code-td44108.html。你可以使用os.environ["OMP_NUM_THREADS"] = "1"等命令，但是必须在导入numpy之前设置。显然，numpy仅在导入时检查此设置。 - seth127
3
@seth127，我的看法是你应该将这个作为答案发布，我认为它很有用。 - kηives
显示剩余2条评论
64

还有更多的环境变量未提及。以下是完整的环境变量列表和使用该变量控制线程数的软件包。请注意，在执行 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
在某些集群上，我发现需要在shell中进行导出，并且os.environ无法工作。 - Shane P Kelly
1
@Amir，这些环境变量的numpy库已经更新了吗？ - zwep
5
你好，这个有任何更新吗？ - a06e
32

关于如何在Python脚本中执行此操作而不是在bash提示符下按照这个帖子，您可以执行以下操作（与上面答案相同）：

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

但是在导入numpy库之前，你必须要将其放在前面。显然，numpy只在导入时检查此项。

5
如果你使用了OpenBLAS，我相信还有一个OPENBLAS_NUM_THREADS选项。 - dga
实际上，这是在你第一次调用MKL优化的Numpy代码时设置，而不是在导入numpy时设置。例如，np.ones()不会触发读取这些变量，但np.dot()会。 - Maxim Imakaev
18
在尝试了上述几种解决方案都没有成功之后，我在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())
有没有办法在笔记本或脚本的开头设置一次？我不想改变我的代码库来包装其中的部分上下文管理器，也不想尝试在这个with语句下包装整个脚本（例如无法处理跨多个单元格的iPython笔记本）。 - MRule
@MRule，如果您采用这种方式（请参见已接受的答案），在执行脚本之前最好将所有环境变量都设置为“1”。 - Gerard
4
我能够在运行时以以下方式解决这个问题：
import mkl
mkl.set_num_threads(1)

我使用以下代码来减少脚本/包中出现问题的可能性：
try:
    import mkl
    mkl.set_num_threads(1)
except:
    pass
网页内容由stack overflow 提供, 点击上面的
可以查看英文原文，
原文链接