NVIDIA 在 scikit-learn、UMAP 和 HDBSCAN 中实现 GPU 加速

Q: 如何在我的 scikit-learn 脚本中启用 cuML 的 GPU 加速？

要启用 cuML 加速，在 Jupyter Notebook 中使用 %load_ext cuml.accel ，或使用 python -m cuml.accel unchanged_script.py 运行脚本。该方式会激活 cuML 的兼容层，将可加速的计算重定向至 GPU。

NVIDIA 的 cuML 是 RAPIDS AI 生态的一部分，是一个基于 GPU 加速的机器学习库。而现在，在开放测试阶段，cuML 25.02 已使您无需修改 Python 代码，就能在 scikit-learn、UMAP 和 HDBSCAN 中启用 GPU 加速。这对期待提升模型性能的 Python 机器学习工程师、数据与科研人员、开发者等众多专业人士来说，都是非常值得关注的消息。

即使您还不熟悉 cuML，甚至不太了解 GPU 加速的概念，也无需担心。本文将带您快速了解一番，读完后，您就能在回归、分类、降维或聚类问题中使用 GPU 加速。

NVIDIA、RAPIDS AI 与 cuML

先简单理清几个关键概念。

什么是 RAPIDS AI？

NVIDIA 开发了为高性能计算、AI 和数据科学提供动力的前沿技术。您或许还记得，当Deepseek 发布 Deepseek-R1 时，人们纷纷讨论其是否使用了 NVIDIA 芯片。

如今，NVIDIA 再次成为焦点，这次是因为其在 RAPIDS AI 生态中的创新。如果您还不熟悉，RAPIDS AI 是一个开源的 GPU 加速库套件，旨在加速数据科学与机器学习工作流。RAPIDS 基于 CUDA（统一计算设备架构），这是一种由 NVIDIA 开发的并行计算平台与 API 模型。RAPIDS 的 GPU 加速库包括：

用于高速 DataFrame 操作的 cuDF
cuML，机器学习库（本文将重点讨论）
cuGraph，图分析库
cuSpatial，地理空间分析库

什么是 cuML？

关于 cuML：cuML 提供了高度优化的经典机器学习算法实现。由于 NVIDIA 拥有业界领先的 AI 芯片，这些在 NVIDIA GPU 上运行的算法能够获得显著的加速效果。与CPU 不同，GPU 擅长并行处理，实际意义在于它们可以同时处理成千上万次计算，非常适合（没错，您猜对了）深度学习和机器学习等任务。

什么是 scikit-learn？

如果您还不熟悉 scikit-learn，同时又认真投入于机器学习与数据科学，强烈建议您报名我们的《使用 scikit-learn 的监督学习》课程。之所以这样说，是因为在机器学习最流行的编程语言中，scikit-learn 是用于表格数据的首选库。此外，scikit-learn 相对易学、API 简洁，并且与 pandas 和 NumPy 配合良好。因此，RAPIDS AI 团队将 scikit-learn 算法作为战略性且自然的重点方向并不令人意外。

cuML 为 scikit-learn 带来了什么

NVIDIA 与 cuML 现已为 scikit-learn、UMAP 和 HDBSCAN 带来一项重大创新：性能速度的巨大提升。根据 RAPIDS AI 团队的说法，使用 scikit-learn 可实现最高 50 倍加速，UMAP 可达 60 倍，HDBSCAN 可达 175 倍。

粗略一算，一个在 scikit-learn 上运行可能需要 5 分钟的算法，理论上可能缩短到 6 秒。这是个很大的差距。下文我们在查看基准测试时会再多谈一些。

零代码改动

如果这一切听起来有点复杂，也不必担心。RAPIDS AI 的目标是：对 scikit-learn 的 API 代码无需任何修改。实际上，您只需要像这样加载扩展：

%load_ext cuml.accel
$ python -m cuml.accel script.py技术支持

这对您熟悉的工作流来说并不是很大的变化——这正是其设计初衷。偶尔您可能需要重启内核并重新加载扩展。

在我看来，“零代码改动”的主张乍听之下确实令人兴奋。其理念在于，您无需修改已经编写好的 scikit-learn 脚本——不只是第一次，以后也不需要。cuML 会在 NVIDIA GPU 上自动加速兼容的组件；如果代码中的某些内容与 NVIDIA GPU 不兼容，脚本将回退到“常规”的 CPU 执行。作为编码者，您只需查看结果，无需操心错误处理或其他问题。

尽管如此，我仍然会预期会有一些边界情况需要调整。在编码中，总会有出错的可能，事情也从来没有看起来那么简单。参考文档中列出了一些内容，多少与“零代码改动”的说法相左，或者说会带来一些摩擦。例如，您需要将列表转换为 NumPy 数组或 pandas DataFrame；字符串标签不受支持，因此需要对类别标签进行预编码。这些也许本来就是常规步骤，但重点是，除了加载扩展之外，您可能还要习惯一些其他事项。

显著的速度提升

如果您只是通过在 mtcars 数据集上运行一个简单的线性回归来练习基础统计，那么为了让斜率和截距快那么一点点而使用 GPU 加速并没有什么意义。省下的零点几秒甚至不值得加载扩展。但如果您运行的是更复杂的算法，节省下的时间就非常重要，主要有两点原因。

首先，节省的时间对许多工作流而言是实实在在的。我提到过最高 50 倍的加速，理论上可以把 5 分钟的训练时间降到 6 秒，节省近 5 分钟。但考虑到机器学习工程师的工作流，模型可能需要多次运行才能调整到位。累加起来，每次运行都节省时间，就意味着数据科学家的整体产出更高。

其次，也是相关的一点，科学家或工程师现在可以尝试更复杂的模型。在前面的例子中，我们假设的是在保持精度水平不变的前提下提升速度。但如果您认为更复杂的模型能带来更高精度呢？想象一下对不同模型参数进行网格搜索并训练模型。（我会在下文尝试一个小示例。）排列组合呈指数增长，使得在 CPU 上的计算成本高得难以承受。但如果您获得 50 倍的加速，这就不再是大问题了。

这些好处是相辅相成的。若项目设置得当，您可能会得到更高精度且训练更快的模型；或者从机器学习工程的角度看，您可以采用更复杂的模型架构，而不会导致过长的训练时间。

数值等效的结果

NVIDIA 特别强调 GPU 将产生数值等效的结果。这是个好消息，意味着加速并不会以牺牲模型精度为代价。

但这也让我想到：为什么说是数值等效而不是完全相同？很可能由于并行处理会带来输出上的细微差异，或者某些方法在数值稳定性上有所不同。比如在 UMAP 的情况下，投影可能会略有差异。我之所以提到这一点，是为了让您在比较 CPU 与 GPU 运行的输出时，若见到细小差别也不必意外。NVIDIA 似乎暗示您应该预期会见到这种情况，而且这是正常的。

当然，如果您挑战更复杂的模型，必须在 GPU 上运行（因为在 CPU 上几乎不可行），那么将其输出与 CPU 等价实现进行比较本身就会变得困难。

cuML 支持的算法与库

前文提到，在本次 cuML 25.02 发布版中已支持 scikit-learn、UMAP 与 HDBSCAN。但熟悉 scikit-learn 的您会知道，其中包含众多不同算法。目前 RAPIDS AI 支持部分库，但并非全部。根据团队介绍，以下是主要支持的方向：

在这些算法中，我对 KNN 的加速最为期待，因为它需要计算所有数据点之间的距离，计算量极大。

下表列出了主要或全部由 cuML 运行的估计器。

算法	库 / 函数名称	类别
UMAP	`umap.UMAP`	降维
PCA（主成分分析）	`sklearn.decomposition.PCA`	降维
截断 SVD	`sklearn.decomposition.TruncatedSVD`	降维
t-SNE	`sklearn.manifold.TSNE`	降维
HDBSCAN	`hdbscan.HDBSCAN`	聚类
K-Means	`sklearn.cluster.KMeans`	聚类
DBSCAN	`sklearn.cluster.DBSCAN`	聚类
随机森林分类器	`sklearn.ensemble.RandomForestClassifier`	分类
逻辑回归	`sklearn.linear_model.LogisticRegression`	分类
K 近邻分类器	`sklearn.neighbors.KNeighborsClassifier`	分类
随机森林回归器	`sklearn.ensemble.RandomForestRegressor`	回归
线性回归	`sklearn.linear_model.LinearRegression`	回归
Elastic Net	`sklearn.linear_model.ElasticNet`	回归
岭回归	`sklearn.linear_model.Ridge`	回归
Lasso 回归	`sklearn.linear_model.Lasso`	回归
核岭回归	`sklearn.kernel_ridge.KernelRidge`	回归
K 近邻回归器	`sklearn.neighbors.KNeighborsRegressor`	回归
最近邻搜索	`sklearn.neighbors.NearestNeighbors`	最近邻搜索

完整的支持方法列表见 RAPIDS AI 文档。预计很快会添加更多算法。（您也可以反馈意见，帮助团队确定优先级。）我注意到的一些缺失项：支持向量回归（sklearn.svm.SVR）、Theil-Sen 回归器（sklearn.linear_model.TheilSenRegressor）、均值漂移聚类（sklearn.cluster.MeanShift）以及多维尺度分析（sklearn.manifold.MDS）。

cuML 的底层工作机制

您可能会好奇其具体如何运作。粗略来看，越来越多的现代库被设计为能自动检测并利用最佳可用的硬件，无论是 CPU 还是 GPU。我们现在看到，当在 scikit-learn 中调用某个函数（启用 cuml.accel 时），软件会检查是否有兼容的 GPU 可用，并在可用时将执行重定向到 GPU 加速版本。

RAPIDS AI 将使 cuml.accel 表现得像 sklearn 的这一机制称为兼容层。在这种情况下，cuML 的兼容层允许 scikit-learn 代码无需修改即可在 NVIDIA GPU 上运行。它被构建为一个即插即用的替代层，充当代理，拦截对 scikit-learn 的函数调用并将其重定向至 cuML。

NVIDIA 的 RAPIDS AI 团队分享了一张很有帮助的图示，我也在此转引：

cuML scikit-learn 工作流示意图。来源：NVIDIA

如何使用 cuML 进行 GPU 加速

cuML 已在 Google Colab 预安装；如果您使用 Jupyter Notebook，在导入 scikit-learn 之前先用以下命令启用：

%load_ext cuml.accel
import sklearn技术支持

接下来，使用 sklearn 的 Python 代码看起来会非常熟悉——这正是重点。以下是一个OLS 回归示例。

# Import necessary libraries
from sklearn.datasets import make_regression
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression

# Generate synthetic regression data
X, y = make_regression(n_samples=500000, n_features=50, noise=0.1, random_state=0)

# Split into training and test sets
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=0)

# Create and train an OLS regression model
ols = LinearRegression()
ols.fit(X_train, y_train)技术支持

如果我想查看哪些部分在 GPU 而非 CPU 上运行，可以使用 logger。

%load_ext cuml.accel

from cuml.common import logger;
logger.set_level(logger.level_enum.debug)技术支持

cuML: Installed accelerator for sklearn.
cuML: Successfully initialized accelerator.
cuML: Performing fit in GPU技术支持

思路就是这样。现在我们来试一个更有说服力的示例。线性回归本身并不太适合并行化，因为求解正规方程或QR 分解涉及一些顺序性操作。

较 OLS 更复杂的一种扩展是岭回归。在使用Python 实现岭回归时，调用 Ridge() 会应用默认超参数，但这未必最优。alpha 值控制正则化强度。过小会过拟合；过大会欠拟合。

%load_ext cuml.accel
import sklearn

from sklearn.linear_model import Ridge
from sklearn.datasets import make_regression
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error

# Generate synthetic data
X, y = make_regression(n_samples=100000, n_features=50, noise=0.1, random_state=0)

# Split into training and test sets
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=0)

# Train a Ridge Regression model with default settings
ridge_default = Ridge()
ridge_default.fit(X_train, y_train)

# Make predictions and evaluate
y_pred_default = ridge_default.predict(X_test)
mse_default = mean_squared_error(y_test, y_pred_default)
技术支持

上述代码仅训练了一个模型，使用单一的 alpha 值和求解器（我认为默认是 'auto'）。虽然可行，但并未针对正则化强度或求解器做优化。

我们可以不依赖默认值，尝试别的方法。可以创建一个网格搜索，测试多组超参数组合并选出最佳者。

from sklearn.model_selection import GridSearchCV

# Define Ridge Regression model
ridge = Ridge()

# Define a grid of hyperparameters to search
param_grid = {
    'alpha': [0.01, 0.1, 1.0, 10.0, 100.0],  # Different regularization strengths
    'solver': ['auto', 'svd', 'cholesky', 'lsqr', 'saga']  # Test multiple solvers
}

# Perform Grid Search with cross-validation
grid_search = GridSearchCV(ridge, param_grid, scoring='neg_mean_squared_error', cv=2, n_jobs=-1)
grid_search.fit(X_train, y_train)

# Get the best model from the search
best_ridge = grid_search.best_estimator_

# Make predictions and evaluate
y_pred_best = best_ridge.predict(X_test)
mse_best = mean_squared_error(y_test, y_pred_best)
技术支持

这里的网格搜索会测试所有可能的超参数组合。若我们为 alpha 定义 5 个取值（[0.01, 0.1, 1.0, 10.0, 100.0]），5 个求解器选项（['auto', 'svd', 'cholesky', 'lsqr', 'saga']），并采用 2 折交叉验证，那么所需训练的模型总数为 5×5×2，即 50 个。这意味着我们不是训练一个模型，而是训练 50 个模型。复杂度大大提升，这正是 GPU 发挥优势的地方——它们要快得多。由于 sklearn（最高）可获得 50 倍加速，训练这 50 个模型所需时间也许与训练一个模型相当。（我有意选择这些数字来说明这个观点。当然，并非所有模型的加速倍数都相同。）

cuML sklearn 基准测试

NVIDIA 比较了在 Intel CPU 与 NVIDIA GPU 上运行的 scikit-learn 算法的性能，涵盖包括分类、回归在内的不同机器学习工作负载。

结论是：像随机森林这样的简单模型，训练时间从 CPU 上的数分钟降至 GPU 上的数秒；而更复杂的聚类模型，则从数小时降至数分钟。所有情形下工作流都更快，且对于高维数据的重计算模型，收益更加明显。这正是 GPU 的强项，因为并行处理可在需要大量矩阵运算、距离计算或以迭代方式进行优化的模型中带来巨大差异。

cuML 下的 scikit-learn（GPU）与 scikit-learn（CPU）对比。来源：NVIDIA

在前面的代码示例中我没有继续讨论线性回归，是想着换个示例更能体现性能。但在这里，我看到线性回归也能获得 52 倍加速。老实说这让我有些意外，因为这个倍数不小。这也提醒我们，NVIDIA 还做了很多高级优化，比如确保不同步骤交错进行，避免 GPU 空转。

一些最佳实践

得益于 cuML 25.02 版，在 GPU 上运行 scikit-learn 算法并不难。正如前文所示，您只需加载扩展。不过，如果您更深入地构建工作流并希望充分利用 GPU 加速，就需要考虑一些更具体的事项。

NVIDIA 建议在流水线中尽量减少 CPU 与 GPU 之间的数据传输。换言之，尽量在 GPU 上完成预处理、训练与推理，然后再把结果送回主机内存。不要把步骤拆散，诸如在 CPU 上做预处理、在 GPU 上做推理这类“混搭”应尽量避免。

此外，为了最大化效率，您还应考虑使用 CUDA-X，这是 NVIDIA 开发的一套 GPU 加速库集合。以随机森林为例，推理阶段可以考虑使用 cuML 的森林推理库，而非 scikit-learn 的实现。

已知限制

RAPIDS AI 团队所做的事情非常复杂，难免会有一些问题。庆幸的是，已知限制的文档维护得很及时。我不打算在这里罗列全部或部分限制，因为很多问题都与具体算法的实现细节相关，而且更新与修补也在快速推进。相反，我给出几个更通用的关注类别。

1. 通用限制

存在一些适用于整个 cuML 的限制，而不仅仅是特定算法。例如，对数据输入格式的限制（与Python 列表有关的问题）、版本兼容性以及内存管理等。

2. 算法特定的限制

某些机器学习算法存在各自独特的约束。这些差异可能影响模型的训练方式、参数行为或结果计算。例如，cuML 中的随机森林在选择分裂阈值的方法上有所不同，这会导致树结构略有差异。

3. 已支持与未支持的功能

在参数、求解器与初始化方法上存在一些差异。例如，PCA 支持 “full” 与 “auto” 的 SVD 求解器，但不支持 “randomized”。KNN 支持闵可夫斯基距离，但不支持马氏距离。

4. 可复现性与数值差异

有时由于并行化或求解器差异，结果可能略有不同。前文也提到过。例如，UMAP 的嵌入结果未必完全一致，但可信度分数应仍然较高。又如，PCA 的成分符号可能会翻转，需要进行归一化。这类都是具体问题，运行特定算法时应参考最新文档。

结语

我对这些在最常用的 Python 机器学习库中、对知名算法实现“易用”的 GPU 加速感到非常期待。这将显著提升模型精度，同时也便捷易上手。

紧跟机器学习前沿。参加我们的《Python 机器学习科学家》职业路径，系统掌握恰当的方法与技术，以便更好地利用最新进展。同时，订阅我们的 DataFramed 播客。我们还有一集与NVIDIA 杰出工程师 Jean-Francois Puget 与 NVIDIA 高级数据科学家 Chris Deotte 的对谈，讨论了包括 GPU 加速影响在内的诸多话题。

作者

Josef Waples

我是数据科学编辑，曾为科学期刊发表的研究论文撰稿。我的兴趣领域包括线性代数、统计学、R 等相关主题。

常见问题

NVIDIA cuML 针对 scikit-learn 的“零代码改动”能力是什么？

NVIDIA cuML 的零代码改动能力，使 scikit-learn 用户无需修改现有代码即可在 NVIDIA GPU 上加速其机器学习工作流。它会自动在 GPU 上运行已支持的算法，对未支持的操作则回退到 CPU。

cuML 的“零代码改动”功能支持哪些 scikit-learn 算法？

该测试版加速了诸如随机森林、k 近邻（KNN）、主成分分析（PCA）和 k-means 聚类等常用的 scikit-learn 算法。后续还将根据用户反馈添加更多算法。

在 NVIDIA GPU 上使用 cuML 可以预期获得多大速度提升？

基准测试显示，加速效果显著：在 NVIDIA GPU 上运行时，scikit-learn 模型最高可快 50 倍，UMAP 可达 60 倍，HDBSCAN 可达 175 倍，相比 CPU 有明显优势。

如何在我的 scikit-learn 脚本中启用 cuML 的 GPU 加速？

要启用 cuML 加速，在 Jupyter Notebook 中使用 %load_ext cuml.accel，或使用 python -m cuml.accel unchanged_script.py 运行脚本。该方式会激活 cuML 的兼容层，将可加速的计算重定向至 GPU。

cuML 能与 pandas、NumPy 等现有 Python 库无缝配合吗？

是的，cuML 旨在与 pandas 和 NumPy 等主流 Python 库协同工作。将 cuDF-pandas 与 cuML 搭配使用，有助于最小化 CPU 与 GPU 之间的数据传输，从而获得更佳性能。

主题

机器学习 Python