Numpy中如何设置OpenBLAS的线程数

在本文中，我们将介绍如何在Numpy中设置OpenBLAS的线程数。OpenBLAS是一种高度优化的线性代数库，可以加速矩阵计算等科学计算任务的执行。然而，默认情况下，OpenBLAS会使用所有的CPU核心，这可能会导致其他任务的执行受到影响。因此，控制OpenBLAS的线程数是很重要的。

阅读更多：Numpy 教程

查看默认的线程数

在开始设置线程数之前，我们需要知道默认情况下OpenBLAS使用了多少线程。我们可以使用numpy.show_config()函数来查看OpenBLAS的设置。

import numpy as np
np.show_config()

运行上面的代码，我们会得到类似下面的输出：

mkl_info:
    libraries = ['mkl_rt', 'pthread']
    library_dirs = ['/opt/conda/lib']
    define_macros = [('SCIPY_MKL_H', None), ('HAVE_CBLAS', None)]
    include_dirs = ['/opt/conda/include']
blas_mkl_info:
    libraries = ['mkl_rt', 'pthread']
    library_dirs = ['/opt/conda/lib']
    define_macros = [('SCIPY_MKL_H', None), ('HAVE_CBLAS', None)]
    include_dirs = ['/opt/conda/include']
blas_opt_info:
    libraries = ['blas', 'blas']
    library_dirs = ['/opt/conda/lib']
    define_macros = [('HAVE_CBLAS', None)]
    include_dirs = ['/opt/conda/include']
lapack_mkl_info:
    libraries = ['mkl_rt', 'pthread']
    library_dirs = ['/opt/conda/lib']
    define_macros = [('SCIPY_MKL_H', None), ('HAVE_CBLAS', None)]
    include_dirs = ['/opt/conda/include']
lapack_opt_info:
    libraries = ['lapack', 'blas', 'blas']
    library_dirs = ['/opt/conda/lib']
    define_macros = [('HAVE_CBLAS', None)]
    include_dirs = ['/opt/conda/include']
openblas_info:
    libraries = ['openblas', 'openblas']
    library_dirs = ['/opt/conda/lib']
    language = c
    define_macros = [('HAVE_CBLAS', None)]
    include_dirs = ['/opt/conda/include']
openblas_lapack_info:
    libraries = ['openblas', 'openblas']
    library_dirs = ['/opt/conda/lib']
    language = c
    define_macros = [('HAVE_CBLAS', None)]
    include_dirs = ['/opt/conda/include']
blas_info:
    libraries = ['blas', 'blas']
    library_dirs = ['/opt/conda/lib']
    define_macros = [('HAVE_CBLAS', None)]
    include_dirs = ['/opt/conda/include']
lapack_info:
    libraries = ['lapack', 'blas', 'blas']
    library_dirs = ['/opt/conda/lib']
    define_macros = [('HAVE_CBLAS', None)]
    include_dirs = ['/opt/conda/include']
lapack_ilp64_info:
    NOT AVAILABLE
blas_ilp64_info:
    NOT AVAILABLE

我们可以看到，上面的输出中包含了很多关于blas（Basic Linear Algebra Subprograms，基本线性代数子程序）和lapack（Linear Algebra Package，线性代数工具包）的信息，这些都是与OpenBLAS相关的库。

我们需要找到openblas_info这个部分，这个部分给出了当前系统上OpenBLAS的设置。在这个例子中，我们可以看到OpenBLAS使用了两个线程（”libraries = [‘openblas’, ‘openblas’]”）。

设置线程数

现在我们已经知道了OpenBLAS的默认线程数，接下来就是如何设置线程数。在Numpy 1.11.0的版本中，Numpy引入了一个环境变量OPENBLAS_NUM_THREADS，用于控制OpenBLAS的线程数。

我们可以在代码中使用os.environ来设置这个环境变量。下面的例子里，我们将设置线程数为1。

import os
import numpy as np

os.environ["OPENBLAS_NUM_THREADS"] = "1"
a = np.random.rand(1000, 1000)
b = np.random.rand(1000, 1000)
c = np.dot(a, b)

在上面的代码中，我们将OPENBLAS_NUM_THREADS设置为1，然后创建两个随机矩阵a和b，并使用numpy.dot()函数计算它们的点积。这个点积的计算使用了OpenBLAS库，而由于我们把线程数设置为1，所以这个计算只会使用一个CPU核心。

检查结果

最后，我们需要检查一下用于计算点积的时间以及结果是否正确。下面的代码将会计算点积和用于计算的时间，并检查一下结果是否与预期相符。

import time

t1 = time.time()
d = np.dot(a, b)
t2 = time.time()

print("Time: {} seconds".format(t2 - t1))
expected = np.dot(a, b)
assert np.allclose(d, expected)

在上面的代码中，我们计算点积及其时间，然后检查结果是否符合我们预期的结果。我们使用numpy.allclose()函数检查两个矩阵是否非常接近。

如果我们再次运行上面这个程序，我们会得到这样的输出：