Question

我正在使用Anaconda的科学Python发行版，其中numpy是用mkl库编译的。

尝试多处理时，我想限制numpy使用单个核心。

我知道我可以设置环境变量＆＃39; MKL_NUM_THREADS＆＃39;，＃NUM; NUMEXPR_NUM_THREADS＆＃39;和＆＃39; OMP_NUM_THREADS＆＃39;到＆＃39; 1＆＃39; （使用os.environ）。但是，这将导致numpy在我的代码周围使用单个核心，并且我希望能够动态地控制它。

我遇到过Python的＆＃39; mkl＆＃39;包（作为Anaconda的一部分分发），其中函数＆＃34; set_num_threads（n）＆＃34;存在。我尝试了它，但它导致了奇怪的行为 - 对于线性代数运算，似乎足够大的矩阵导致numpy恢复多线程。

任何人都可以提出有关动态控制numpy / mkl多线程的解决方案/主管/任何内容吗？

非常感谢！

Answer 1

要从 python 动态控制 MKL 或其他 BLAS 库（OpenBLAS、Blis）线程，您可以使用 threadpoolctl 模块。

使用pip安装

pip install threadpoolctl

例如：-

考虑到您的 NumPy 安装与某些 BLAS 库（MKL、OpenBLAS、Blis）相关联。这将为您的 BLAS 运行时动态地将线程设置为 1。

import threadpoolctl
import numpy as np

A = np.random.random((10**4))
B = np.random.random((10**4))

with threadpoolctl.threadpool_limits(limits=1, user_api="blas"):
    # MKL threads will be limited to 1 inside this block
    np.dot(A, B)

进一步参考：https://github.com/joblib/threadpoolctl#usage

动态控制基于mkl的numpy多线程

1 个答案: