Scikit-Learn 1.4使用指南：数组 API 支持（实验性）Array API support (experimental)

scikit-learn 1.6 官网中文文档翻译已经上线了，欢迎使用：http://www.aidoczh.com/scikit-learn/

sklearn

数组 API 规范定义了一个具有类似于 NumPy API 的标准 API，用于所有具有数组操作功能的库。Scikit-learn 的数组 API 支持需要安装 array-api-compat。

一些主要依赖于 NumPy（而不是使用 Cython）来实现其拟合、预测或转换方法的 scikit-learn 估计器可以配置为接受任何与数组 API 兼容的输入数据结构，并自动将操作分派到底层命名空间，而不是依赖于 NumPy。

目前，此支持被视为实验性的，必须按照以下说明显式启用。

[!NOTE]
目前，只有 cupy.array_api、numpy.array_api、cupy 和 PyTorch 被认为与 scikit-learn 的估计器兼容。

示例用法

以下是一个示例代码片段，演示如何使用 CuPy 在 GPU 上运行 ~discriminant_analysis.LinearDiscriminantAnalysis：

from sklearn.datasets import make_classification
from sklearn import config_context
from sklearn.discriminant_analysis import LinearDiscriminantAnalysis
import cupy

X_np, y_np = make_classification(random_state=0)
X_cu = cupy.asarray(X_np)
y_cu = cupy.asarray(y_np)
X_cu.device
# <CUDA Device 0>

with config_context(array_api_dispatch=True):
    lda = LinearDiscriminantAnalysis()
    X_trans = lda.fit_transform(X_cu, y_cu)
X_trans.device
# <CUDA Device 0>

模型训练后，数组类型的拟合属性也将与训练数据相同，属于相同的数组 API 命名空间。例如，如果使用了 CuPy 的数组 API 命名空间进行训练，则拟合属性将位于 GPU 上。我们提供了一个实验性的 _estimator_with_converted_arrays 实用程序，用于将估计器属性从数组 API 转换为 ndarray：

from sklearn.utils._array_api import _estimator_with_converted_arrays
cupy_to_ndarray = lambda array : array.get()
lda_np = _estimator_with_converted_arrays(lda, cupy_to_ndarray)
X_trans = lda_np.transform(X_np)
type(X_trans)
# <class 'numpy.ndarray'>

PyTorch 支持

通过设置 array_api_dispatch=True 并直接传入张量，支持 PyTorch 张量：

import torch
X_torch = torch.asarray(X_np, device="cuda", dtype=torch.float32)
y_torch = torch.asarray(y_np, device="cuda", dtype=torch.float32)

with config_context(array_api_dispatch=True):
    lda = LinearDiscriminantAnalysis()
    X_trans = lda.fit_transform(X_torch, y_torch)
type(X_trans)
# <class 'torch.Tensor'>
X_trans.device.type
# 'cuda'

支持数组 API 兼容输入

scikit-learn 中支持数组 API 兼容输入的估计器和其他工具。

估计器

• decomposition.PCA（使用 svd_solver=“full”、svd_solver=“randomized” 和 power_iteration_normalizer=“QR”）
• discriminant_analysis.LinearDiscriminantAnalysis（使用 solver=“svd”）
• preprocessing.KernelCenterer
• preprocessing.MaxAbsScaler
• preprocessing.MinMaxScaler
• preprocessing.Normalizer

评估指标

• sklearn.metrics.accuracy_score
• sklearn.metrics.zero_one_loss

工具

• model_selection.train_test_split

预计覆盖范围将随时间增加。请关注专用的 GitHub 上的元问题 以跟踪进展。

常见估计器检查

将 array_api_support 标签添加到估计器的标签集中，以指示其支持数组 API。这将启用专用检查作为常规测试的一部分，以验证在使用原始 NumPy 和数组 API 输入时估计器的结果是否相同。

要运行这些检查，您需要在测试环境中安装 array_api_compat。要运行完整的检查集，您需要同时安装 PyTorch 和 CuPy，并且需要有一个 GPU。无法执行或存在缺失依赖项的检查将自动跳过。因此，重要的是使用 -v 标志运行测试，以查看跳过了哪些检查：

pip install array-api-compat  # and other libraries as needed
pytest -k "array_api" -v

MPS 设备支持说明

在 macOS 上，PyTorch 可以使用 Metal Performance Shaders（MPS）访问硬件加速器（例如 M1 或 M2 芯片的内部 GPU 组件）。然而，目前编写本文时，PyTorch 对 MPS 设备的支持不完整。有关更多详细信息，请参阅以下 GitHub 问题：

• https://github.com/pytorch/pytorch/issues/77764

要在 PyTorch 中启用 MPS 支持，请在运行测试之前设置环境变量 PYTORCH_ENABLE_MPS_FALLBACK=1：

PYTORCH_ENABLE_MPS_FALLBACK=1 pytest -k "array_api" -v

目前，所有 scikit-learn 的测试应该通过，但计算速度不一定比使用 CPU 设备更快。