Scikit-Learn 1.8引入 Array API,支持 PyTorch 与 CuPy 张量的原生 GPU 加速

Scikit-Learn 1.8.0 更新引入了实验性的 Array API 支持。这意味着 CuPy 数组或 PyTorch 张量现在可以直接在 Scikit-Learn 的部分组件中直接使用了,且计算过程能保留在 GPU 上。

1.8.0 到底更新了什么?

Scikit-Learn 开始正式支持Python Array API 标准。这是一个由 NumPy、CuPy、PyTorch、JAX 等库共同维护的接口规范。在 1.8.0 版本中可以实现:

  • 直接传参:受支持的评估器(estimators)现在可以直接接收 CuPy 数组或 PyTorch 张量。
  • 计算分派:运算会被自动分派到对应的非 CPU 设备(如 GPU)上执行。
  • 状态保留:模型拟合后的属性会与输入数据保持在同一物理设备上。

虽然目前的版本依然贴着"实验性"标签且需要显式开启,但它确实打破了 Scikit-Learn 过去那种"万物皆需 NumPy"的框架。

交叉验证

如果你平时不怎么用

复制代码
cross_val_score

复制代码
GridSearchCV

复制代码
CalibratedClassifierCV

,那你可能感觉不到这次更新的提速。但对大多数从事肃建模的开发者来说,交叉验证一直是 GPU 的"性能杀手"。

在旧版本中,即便你的基础模型(如 XGBoost)是在 GPU 上训练的,Scikit-Learn 的编排逻辑会把数组转回 NumPy,然后在 CPU 上重新计算各项指标。这种频繁的内存搬运和 CPU 的操作浪费了大量的时间,但是Array API 的加入让这种循环能基本闭环在 GPU 内部运行。

开启方式与限制

启用这项特性需要完成下面的配置。如果漏掉任何一步,程序都会悄悄退回到 NumPy 模式。

环境变量设置(必须在导入 SciPy 或 Scikit-Learn 之前):

复制代码
 importos  
 os.environ["SCIPY_ARRAY_API"] ="1"

配置 Scikit-Learn 内部开关

复制代码
 fromsklearnimportset_config  
 set_config(array_api_dispatch=True)

目前还有一个问题,就是不支持 cuDF DataFrames。但是你依然可以用 cuDF 做数据加载和预处理,不过输入模型之前必须确保输入是 array-like 格式。也就是说类别特征必须手动编码而且且无法再依赖 pandas/cuDF 的 dtype 自动识别机制。

基于 GPU 的 XGBoost 交叉验证

下面是一个运行 5 折分层交叉验证的示例。为了让整个链路留在 GPU 上,我们需要对

复制代码
XGBClassifier

做一点小的封装,并结合 cuML 的指标计算。

复制代码
 import os  
 os.environ['SCIPY_ARRAY_API'] = '1'  
   
 import cupy as cp  
 import cudf  
 from sklearn.model_selection import StratifiedKFold, cross_val_score  
 from sklearn.metrics import make_scorer  
 from cuml.metrics import roc_auc_score  
 from xgboost import XGBClassifier  
 from sklearn import set_config  
 set_config(array_api_dispatch=True)  
 
 # 加载数据并进行简单的预处理
 X = cudf.read_csv('/kaggle/input/playground-series-s5e12/train.csv').set_index('id')  
 y = X.pop('diagnosed_diabetes').astype(int)  
 
 # 类别特征编码处理
 cat_cols = [c for c in X.columns if X[c].dtype == 'object']  
 X = X.astype({c: 'category' for c in cat_cols})  
 for c in cat_cols:  
     X[c] = X[c].cat.codes  
 
 ft = ['c' if c in cat_cols else 'q' for c in X.columns]  
 kfold = StratifiedKFold(5, shuffle=True, random_state=0)  
 
 # 封装 XGB 以适配 CuPy 预测
 class cuXGBClassifier(XGBClassifier):  
     @property  
     def classes_(self):  
         return cp.asarray(super().classes_)  
     def predict_proba(self, X):  
         p = self.get_booster().inplace_predict(X)  
         if p.ndim == 1:  
             p = cp.column_stack([1 - p, p])  
         return p  
     def predict(self, X):  
         return cp.asarray(super().predict(X))  
 
 model = cuXGBClassifier(  
     enable_categorical=True,  
     feature_types=ft,  
     device='cuda',  
     n_jobs=4,  
     random_state=0  
 )  
 
 # 执行交叉验证
 scores = cross_val_score(  
     model,  
     X.values,  
     y.values,  
     cv=kfold,  
     scoring=make_scorer(  
         roc_auc_score,  
         response_method="predict_proba"  
     ),  
     n_jobs=1  
 )  
 print(f"{scores.mean():.5f} ± {scores.std():.5f}")

虽然这段代码看起来还是需要一些修改,但它确实能让交叉验证循环保持在 GPU 上。

现阶段支持的组件

目前 Array API 的覆盖范围还在逐步扩大。在 1.8.0 中,以下组件已经具备了较好的支持:

  • 预处理StandardScalerPolynomialFeatures
  • 线性模型与校准RidgeCVRidgeClassifierCVCalibratedClassifierCV
  • 聚类与混合模型GaussianMixture

官方提供的一个基于 PyTorch 的 Ridge 管道示例显示,在处理线性代数密集型任务时,这种配置在 Colab 环境下能比单核 CPU 快出 10 倍左右。

复制代码
 ridge_pipeline_gpu = make_pipeline(  
     feature_preprocessor,  
     FunctionTransformer(  
         lambda x: torch.tensor(  
             x.to_numpy().astype(np.float32),  
             device="cuda"  
         )  
     ),  
     CalibratedClassifierCV(  
         RidgeClassifierCV(alphas=alphas),  
         method="temperature"  
     ),  
 )  
   
 with sklearn.config_context(array_api_dispatch=True):  
     cv_results = cross_validate(  
         ridge_pipeline_gpu, features, target  
     )

总结

Scikit-Learn 准备好完全接管 GPU 了吗?显然还没有。但这个版本意义在于,它正已经向GPU的支持迈出了第一步。目前这种方式虽然还有点"硬核",对普通用户不够友好,但对于追求极致效率的开发者来说,Scikit-Learn 1.8.0 已经要想这个方向前进了。

Scikit-Learn 1.8.0 首次引入实验性 Array API 支持,可直接使用 CuPy 数组或 PyTorch 张量,计算全程保留在 GPU。交叉验证等操作不再强制转回 CPU,大幅提升效率。需配置环境变量与 set_config 开启,目前支持部分组件如 Ridge、GaussianMixture 等,标志其迈向 GPU 加速的重要一步。

作者:Abish Pius

相关推荐
KG_LLM图谱增强大模型2 小时前
因果知识图谱增强大模型重构生命历程风险路径和慢病防控新范式:从妊娠期糖尿病到痴呆症的研究突破
人工智能·重构·知识图谱
万俟淋曦2 小时前
【论文速递】2025年第37周(Sep-07-13)(Robotics/Embodied AI/LLM)
人工智能·ai·机器人·大模型·论文·robotics·具身智能
free-elcmacom2 小时前
机器学习高阶教程<11>当数据开始“折叠”:流形学习与深度神经网络如何发现世界的隐藏维度
人工智能·python·神经网络·学习·算法·机器学习·dnn
阿杰学AI2 小时前
AI核心知识61——大语言模型之Embedding (简洁且通俗易懂版)
人工智能·机器学习·ai·语言模型·自然语言处理·embedding·词向量
RockHopper20252 小时前
一种面向服务LLM应用系统的显式世界模型架构原理
人工智能·llm·世界模型·显式模型
tap.AI2 小时前
(一)初识 Stable Diffusion 3.5 —— 下一代多模态架构详解
人工智能·stable diffusion
Master_oid2 小时前
机器学习26:增强式学习(Deep Reinforcement Learn)①
人工智能·学习·机器学习
鲨莎分不晴2 小时前
从 0 实现一个 Offline RL 算法 (以 IQL 为例)
人工智能·深度学习·机器学习
rayufo2 小时前
深度学习图像复原论文《SwinIR: Image Restoration Using Swin Transformer》解读及其代码实现
人工智能·深度学习·transformer