「scikit-learn」作为经典的机器学习框架，其运算速度一直广受用户的诟病。今天分享一个新包 「sklearnex」，可以在不改变原有代码的基础上，获得数十倍甚至上千倍的效率提升。


sklearnex简介

借助英特尔® Extension for Scikit-learn，您可以加速您的 Scikit-learn 应用程序，并且不需修改 Scikit-Learn原有代码结构。 英特尔® Extension for Scikit-learn是一款免费软件 AI 加速器，可带来超过 10-100 倍的加速。


安装sklearnex

pip3 install notebook scikit-learn-intelex



快速上手

scikit-learn原代码

import numpy as np
from sklearn.cluster import KMeans

X = np.array([[1,  2], [1,  4], [1,  0],
              [10, 2], [10, 4], [10, 0]])

kmeans = KMeans(n_clusters=2, random_state=0).fit(X)
print(f"kmeans.labels_ = {kmeans.labels_}")

Run

kmeans.labels_ = [1 1 1 0 0 0]

加速代码

import numpy as np
from sklearnex import patch_sklearn
patch_sklearn() #启动加速补丁

# 加速补丁放置于sklearn之前
from sklearn.cluster import KMeans

X = np.array([[1,  2], [1,  4], [1,  0],
              [10, 2], [10, 4], [10, 0]])
kmeans = KMeans(n_clusters=2, random_state=0).fit(X)
print(f"kmeans.labels_ = {kmeans.labels_}")

Run

kmeans.labels_ = [1 1 1 0 0 0]

去掉加速

一般使用sklearn的环境是jupyter内，如果使用了sklearnex的加速功能，那么整个jupyter文件都是加速环境。如果想回到常规速度，可以在机器学习算法之前使用unpatch_sklearn()回到sklearn正常速度

from sklearnex import unpatch_sklearn

unpatch_sklearn()

# Re-import scikit-learn algorithms after the unpatch:
from sklearn.cluster import KMeans



效率对比

按照官方说法，越强劲的CPU可以获得的性能提升比例也会更高，下图是官方在Intel Xeon Platinum 8275CL处理器下测试了一系列算法后得出的性能提升结果，不仅可以提升训练速度，还可以提升模型推理预测速度，在某些场景下甚至达到数千倍的性能提升。



广而告之