用一行代码将 sklearn 运算速度提升上千倍

1. 简介

作为经典的机器学习框架，scikit-learn 已经发展了十余年，但其运算速度一直是用户诟病的问题。尽管 scikit-learn 自带的一些运算加速功能效果有限，但并不能充分利用算力。

今天我要介绍的内容，可以帮助我们在不改变原有代码的基础上，显著提升 scikit-learn 的运算效率，让我们一起来看看吧！

2. 利用 sklearnex 加速 scikit-learn

为了加速运算，我们只需要安装 sklearnex 这个拓展库，就能在拥有 Intel 处理器的设备上获得大幅度的运算效率提升。

首先，我们可以在一个独立的 conda 虚拟环境中进行实验。下面是设置实验环境的全部命令，同时我们也安装 jupyterlab 作为 IDE：

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

准备好实验环境后，我们在 Jupyter Lab 中编写测试代码来看看加速效果。使用起来非常简单，只需在代码中导入 scikit-learn 相关功能模块之前，运行以下代码：

<code>from sklearnex import patch_sklearn, unpatch_sklearn patch_sklearn()</code>

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

其他要做的就是继续执行你原来的 scikit-learn 代码。我在我的老款 Lenovo 拯救者笔记本上进行了简单测试。

以线性回归为例，在百万级别样本量和上百个特征的示例数据集上，开启加速后仅耗时 0.21 秒就完成对训练集的训练；而使用 unpatch_sklearn() 强制关闭加速模式后，训练耗时随即上升到 11.28 秒，意味着通过 sklearnex 我们获得了 50 多倍的运算速度提升！

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

官方还提供了一些 ipynb 示例，展示了包括 K-means、DBSCAN、随机森林、逻辑回归、岭回归等多种常用算法的示例。感兴趣的读者可以自行下载学习。

通过 sklearnex，我们可以用一行代码让 scikit-learn 的运算速度提升上千倍。希望这篇文章对你有所帮助！