如何在Google Colab中禁用Numpy与Scikit-learn的多线程？

问题背景 #

核心需求是使用Python multiprocessing.Pool 基于同一数据并行运行大量sklearn LogisticRegression 实例。对比单线程与多进程代码运行时间时发现：

• 多进程代码并未比单线程更快
• 多进程中print(time() - t1)的输出远高于单线程

原因是Numpy默认使用多线程，导致进程间CPU资源争抢。本地环境通过设置OPENBLAS_NUM_THREADS=1解决，但迁移到Google Colab后，即使设置多个相关环境变量也无效。

复现代码 #

通用函数与数据生成 #

def f(data):
    t1 = time()
    X_tr, X_te, y_tr, y_te = data
    lr = LogisticRegression(n_jobs=1)
    lr.fit(X_tr, y_tr)
    res = accuracy_score(y_te, lr.predict(X_te))
    print(time() - t1)
    return res

K = 20

data = np.random.random((100000, 5))
target = np.random.randint(0, 10, (100000))

单线程代码 #

res = []
for i in range(K):
    res.append(f(train_test_split(data.copy(), target.copy(), train_size=0.7)))
print(sum(res) / len(res))

多进程代码 #

with Pool(10) as p:
    res = list(p.map(f, [train_test_split(data.copy(), target.copy(), train_size=0.7) for x in range(K)]))
print(sum(res) / len(res))

完整测试代码（本地有效，Colab无效） #

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

from multiprocessing import Pool, cpu_count
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score
from time import time
import numpy as np

def f(data):
    t1 = time()
    X_tr, X_te, y_tr, y_te = data
    lr = LogisticRegression(n_jobs=1)
    lr.fit(X_tr, y_tr)
    res = accuracy_score(y_te, lr.predict(X_te))
    print(time() - t1)
    return res

K = 20

data = np.random.random((100000, 5))
target = np.random.randint(0, 10, (100000))

res = []
for i in range(K):
    res.append(f(train_test_split(data.copy(), target.copy(), train_size=0.7)))
print(sum(res) / len(res))

with Pool(10) as p:
    res = list(p.map(f, [train_test_split(data.copy(), target.copy(), train_size=0.7) for x in range(K)]))
print(sum(res) / len(res))

尝试过的无效方法 #

在Python代码中设置以下环境变量，无法解决Colab中的问题：

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

解决方案 #

方法1：使用Colab魔法命令提前设置环境变量 #

在导入任何库之前，运行以下魔法命令。%env会在Python进程启动前配置环境变量，确保numpy加载时读取到单线程设置：

%env OMP_NUM_THREADS=1
%env OPENBLAS_NUM_THREADS=1
%env MKL_NUM_THREADS=1
%env VECLIB_MAXIMUM_THREADS=1
%env NUMEXPR_NUM_THREADS=1

方法2：针对性设置BLAS后端线程数 #

先执行以下代码查看numpy使用的BLAS/LAPACK后端：

import numpy as np
np.__config__.show()

根据输出结果，只设置对应后端的线程数。例如如果显示使用MKL，仅设置MKL_NUM_THREADS=1即可。

方法3：在子进程内部强制设置 #

若上述方法无效，可在任务函数内部重新设置环境变量并重新加载numpy：

def f(data):
    import os
    os.environ["OMP_NUM_THREADS"] = "1"
    os.environ["MKL_NUM_THREADS"] = "1"
    
    import importlib
    import numpy as np
    importlib.reload(np)
    
    t1 = time()
    X_tr, X_te, y_tr, y_te = data
    lr = LogisticRegression(n_jobs=1)
    lr.fit(X_tr, y_tr)
    res = accuracy_score(y_te, lr.predict(X_te))
    print(time() - t1)
    return res

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