scipy 在python中设置线程数

okxuctiv  于 2022-11-10  发布在  Python
关注(0)|答案(1)|浏览(208)

我有一个python脚本,它由numpy和scipy函数组成。我试图检查我的代码的缩放。

numpy.show_config()

安装在我的系统中的numpy的配置显示以下信息。

blas_mkl_info:
    libraries = ['mkl_rt']
    library_dirs = ['Library\\lib']
    define_macros = [('SCIPY_MKL_H', None), ('HAVE_CBLAS', None)]
    include_dirs = ['Library\\include']
blas_opt_info:
    libraries = ['mkl_rt']
    library_dirs = ['Library\\lib']
    define_macros = [('SCIPY_MKL_H', None), ('HAVE_CBLAS', None)]
    include_dirs = ['Library\\include']
lapack_mkl_info:
    libraries = ['mkl_rt']
    library_dirs = ['Library\\lib']
    define_macros = [('SCIPY_MKL_H', None), ('HAVE_CBLAS', None)]
    include_dirs = ['Library\\include']
lapack_opt_info:
    libraries = ['mkl_rt']
    library_dirs = ['Library\\lib']
    define_macros = [('SCIPY_MKL_H', None), ('HAVE_CBLAS', None)]
    include_dirs = ['Library\\include']
Supported SIMD extensions in this NumPy install:
    baseline = SSE,SSE2,SSE3
    found = SSSE3,SSE41,POPCNT,SSE42,AVX,F16C,FMA3,AVX2,AVX512F,AVX512CD,AVX512_SKX,AVX512_CLX,AVX512_CNL
    not found =

因此,我尝试在导入numpy之前设置以下环境变量

import os    
os.environ["OMP_NUM_THREADS"] = '16' 
os.environ["OPENBLAS_NUM_THREADS"] = '16' 
os.environ["MKL_NUM_THREADS"] = '16'

但是,我的代码仍然只使用1个线程,程序的执行时间没有差异。
我也试过设置mkl.set_num_threads(16),但没有区别。
我知道python有GIL,它不允许多个线程同时执行,就像你在C中期望的那样。有没有其他方法来设置python中使用的线程数?

pdkcd3nj

pdkcd3nj1#

它可以工作,但您必须在第一次加载脚本中的模块(包括任何子模块)之前设置环境变量。
为了安全起见,您应该在任何其他导入之前在主程序中执行此操作

import os
nthreads = 1
os.environ["OMP_NUM_THREADS"] = str(nthreads) 
os.environ["OPENBLAS_NUM_THREADS"] = str(nthreads) 
os.environ["MKL_NUM_THREADS"] = str(nthreads)
import numpy

另一种方法是按照this answer的建议,在运行脚本之前设置环境变量。
有一件事我认为是,使用导入库重新加载模块可以做的技巧,允许设置它的动态,但没有,它不工作。
如果您可以使用pytorch而不是numpy,那么您可以使用torch.set_num_threads,这是有效的,请查看here的示例用法。

相关问题