Linux环境下如何让NumPy启用OpenBLAS实现多线程矩阵乘法加速？

问题背景 #

尝试通过NumPy调用OpenBLAS实现矩阵乘法的CPU并行加速，按文档说明numpy.dot会使用优化的BLAS库，但安装OpenBLAS后，htop显示仅使用8核CPU中的1核，运行速度未得到提升。

环境与安装步骤 #

• 系统：最新版Linux Mint
• 安装操作：
  1. 执行sudo apt-get install libopenblas-dev安装OpenBLAS开发包
  2. 通过pip卸载并重装NumPy，未手动编译OpenBLAS源码

测试脚本 #

import numpy as np
import time
import multiprocessing as mp
import os

#Environment for multi-threading
nb_processeurs = str(mp.cpu_count())
os.environ["OPENBLAS_NUM_THREADS"] = nb_processeurs
os.environ["BLAS"] = "openblas64_"

print(np.__config__.show())

#Variables
n = 5000
p = 300

# Generate a matrix of 0 and 1 with 30% of 1 and 70% of 0
A = (np.random.rand(n,p)> 0.7).astype(int)
A_t = A.T

#Numpy dot product
start_time = time.time()
C1 = np.dot(A,A_t)
end_time = time.time()
print("NumPy dot product took {} seconds".format(round(end_time - start_time,2)))

NumPy配置输出 #

openblas64__info:
    libraries = ['openblas64_', 'openblas64_']
    library_dirs = ['/usr/local/lib']
    language = c
    define_macros = [('HAVE_CBLAS', None), ('BLAS_SYMBOL_SUFFIX', '64_'), ('HAVE_BLAS_ILP64', None)]
    runtime_library_dirs = ['/usr/local/lib']
blas_ilp64_opt_info:
    libraries = ['openblas64_', 'openblas64_']
    library_dirs = ['/usr/local/lib']
    language = c
    define_macros = [('HAVE_CBLAS', None), ('BLAS_SYMBOL_SUFFIX', '64_'), ('HAVE_BLAS_ILP64', None)]
    runtime_library_dirs = ['/usr/local/lib']
openblas64__lapack_info:
    libraries = ['openblas64_', 'openblas64_']
    library_dirs = ['/usr/local/lib']
    language = c
    define_macros = [('HAVE_CBLAS', None), ('BLAS_SYMBOL_SUFFIX', '64_'), ('HAVE_BLAS_ILP64', None), ('HAVE_LAPACKE', None)]
    runtime_library_dirs = ['/usr/local/lib']
lapack_ilp64_opt_info:
    libraries = ['openblas64_', 'openblas64_']
    library_dirs = ['/usr/local/lib']
    language = c
    define_macros = [('HAVE_CBLAS', None), ('BLAS_SYMBOL_SUFFIX', '64_'), ('HAVE_BLAS_ILP64', None), ('HAVE_LAPACKE', None)]
    runtime_library_dirs = ['/usr/local/lib']
Supported SIMD extensions in this NumPy install:
    baseline = SSE,SSE2,SSE3
    found = SSSE3,SSE41,POPCNT,SSE42,AVX,F16C,FMA3,AVX2
    not found = AVX512F,AVX512CD,AVX512_KNL,AVX512_KNM,AVX512_SKX,AVX512_CLX,AVX512_CNL,AVX512_ICL
None

解决方案 #

1. 调整环境变量设置顺序 #

脚本中先导入NumPy再设置环境变量，但NumPy在导入时就已加载BLAS库，后续设置的环境变量无法生效。修改脚本，将环境变量设置放在NumPy导入前：

import os
import multiprocessing as mp

# 先设置环境变量
nb_processeurs = str(mp.cpu_count())
os.environ["OPENBLAS_NUM_THREADS"] = nb_processeurs
os.environ["BLAS"] = "openblas64_"

# 再导入NumPy
import numpy as np
import time

# 后续代码不变
print(np.__config__.show())

n = 5000
p = 300
A = (np.random.rand(n,p)> 0.7).astype(int)
A_t = A.T

start_time = time.time()
C1 = np.dot(A,A_t)
end_time = time.time()
print("NumPy dot product took {} seconds".format(round(end_time - start_time,2)))

2. 系统层面配置OpenBLAS线程数（可选） #

若不想每次在脚本中设置，可配置系统全局环境变量：

• 编辑~/.bashrc文件，添加：

  export OPENBLAS_NUM_THREADS=8  # 替换为你的CPU核心数
  export BLAS=openblas64_
  

• 执行source ~/.bashrc使配置生效

3. 验证NumPy与OpenBLAS的链接 #

执行以下命令，确认NumPy确实依赖OpenBLAS：

ldd $(python -c "import numpy; print(numpy.__file__)") | grep openblas

若输出包含libopenblas64_.so，说明链接正常。

4. 确认OpenBLAS的编译配置 #

执行openblas-config --show-config，查看是否开启多线程支持。若输出中包含USE_OPENMP=1或NUM_THREADS=8，说明多线程功能正常。

内容的提问来源于stack exchange，提问作者Antoine