(数据科学学习手札138)使用sklearnex大幅加速scikit-learn运算

本文示例代码已上传至我的Github仓库//github.com/CNFeffery/DataScienceStudyNotes

1 简介

  大家好我是费老师,scikit-learn作为经典的机器学习框架,从诞生至今已发展了十余年,但其运算速度一直广受用户的诟病。熟悉scikit-learn的朋友应该清楚,scikit-learn中自带的一些基于joblib等库的运算加速功能效果有限,并不能很充分地利用算力。

  而今天我要给大家介绍的知识,可以帮助我们在不改变原有代码的基础上,获得数十倍甚至上千倍的scikit-learn运算效率提升,let’s go!

2 利用sklearnex加速scikit-learn

  为了达到加速运算的效果,我们只需要额外安装sklearnex这个拓展库,就可以帮助我们在拥有intel处理器的设备上,获得大幅度的运算效率提升。

  抱着谨慎尝鲜的态度,我们可以在单独的conda虚拟环境中做实验,全部命令如下,我们顺便安装jupyterlab作为IDE:

conda create -n scikit-learn-intelex-demo python=3.8 -c //mirrors.sjtug.sjtu.edu.cn/anaconda/pkgs/main -y
conda activate scikit-learn-intelex-demo
pip install scikit-learn scikit-learn-intelex jupyterlab -i //pypi.douban.com/simple/

  完成实验环境的准备后,我们在jupyter lab中编写测试用代码来看看加速效果如何,使用方式很简单,我们只需要在代码中导入scikit-learn相关功能模块之前,运行下列代码即可:

from sklearnex import patch_sklearn, unpatch_sklearn
patch_sklearn()

  成功开启加速模式后会打印以下信息:

  其他要做的仅仅是将你原本的scikit-learn代码在后面继续执行即可,我在自己平时写作以及开发开源项目的老款拯救者笔记本上简单测试了一下。

  以线性回归为例,在百万级别样本量以及上百个特征的示例数据集上,开启加速后仅耗时0.21秒就完成对训练集的训练,而使用unpatch_sklearn()强制关闭加速模式后(注意scikit-learn相关模块需要重新导入),训练耗时随即上升到11.28秒,意味着通过sklearnex我们获得了50多倍的运算速度提升!

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

  官方也提供了一些ipynb示例(//github.com/intel/scikit-learn-intelex/tree/master/examples/notebooks),展示了包含K-meansDBSCAN随机森林逻辑回归岭回归等多种常用算法示例,感兴趣的读者朋友们可以自行下载学习。


  以上就是本文的全部内容,欢迎在评论区与我进行讨论~