Scikit-Learn 1.8引入 Array API，支持 PyTorch 與 CuPy 張量的原生 GPU 加速詳情 - 人工智能,深度學習,神經網絡,pytorch,sklearn-pandas deephub 博客

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 中，以下組件已經具備了較好的支持：

預處理：StandardScaler、PolynomialFeatures
線性模型與校準：RidgeCV、RidgeClassifierCV、CalibratedClassifierCV
聚類與混合模型：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 已經要想這個方向前進了。

https://avoid.overfit.cn/post/ab7e632896364fc3b4b9fdc9d17884e3

作者：Abish Pius

deephub 博客

deephub 博客

博客 / 詳情