13.2. Поддержка Array API (экспериментальная)

13.2.1. Поддерживаемые библиотеки для работы с массивами

В следующей таблице перечислены библиотеки и аппаратные платформы, для которых мы регулярно проводим автоматические тесты на соответствие требованиям. Другие библиотеки и устройства, соответствующие Array API, также могут работать без дополнительной настройки.

Ожидается, что со временем охват будет расширяться.

13.2.2. Включение поддержки Array API

Чтобы включить поддержку Array API, параметру конфигурации array_api_dispatch необходимо присвоить значение True. Рекомендуем задать этот параметр глобально, чтобы обеспечить единообразное поведение и предотвратить случайное смешивание пространств имён массивов. Обратите внимание: в примерах ниже мы используем менеджер контекста (config_context), чтобы не сбрасывать значение на False в конце каждого фрагмента кода и тем самым не влиять на остальную документацию.

Для поддержки стандарта Array API в scikit-learn необходимо до импорта scipy и scikit-learn установить переменную окружения SCIPY_ARRAY_API в значение 1:

export SCIPY_ARRAY_API=1

Обратите внимание: эта переменная окружения предназначена для временного использования. Подробнее см. в документации SciPy по Array API.

Для работы с Array API предполагается, что установлены последние версии зависимостей scikit-learn. Более старые версии могут работать, но мы этого не гарантируем. Хотя поддержка Array API отмечена как экспериментальная, обратная совместимость не гарантируется. В частности, если в новой версии зависимости исправлена ошибка, мы не будем добавлять дополнительный код для переноса исправления или поддержания совместимости со старыми версиями.

Scikit-learn принимает входные данные типа array-like для всех metrics и некоторых оценивателей. Если array_api_dispatch=False, эти данные преобразуются в массивы NumPy с помощью numpy.asarray (или numpy.array). Хотя таким способом успешно преобразуются некоторые входные данные Array API (например, массив JAX), при использовании таких входных данных мы обычно рекомендуем задавать array_api_dispatch=True. Это связано с тем, что преобразование в NumPy часто может завершиться ошибкой, например, для тензора torch, размещённого на GPU.

13.2.3. Пример использования

Приведённый ниже фрагмент кода демонстрирует, как использовать PyTorch для запуска LinearDiscriminantAnalysis на GPU CUDA:

>>> from sklearn.datasets import make_classification
>>> from sklearn import config_context
>>> from sklearn.discriminant_analysis import LinearDiscriminantAnalysis
>>> import torch

>>> X_np, y_np = make_classification(random_state=0)
>>> 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'

Этот подход работает одинаково с любой поддерживаемой библиотекой для работы с массивами. Например, замените torch.asarray(..., device="cuda") на cupy.asarray(...) для CuPy или на dpnp.asarray(...) для dpnp. В PyTorch также можно использовать другие устройства, изменив аргумент device= (например, "cpu", "xpu", "mps").

После обучения модели атрибуты обученного оценивателя, содержащие массивы, также будут принадлежать тому же пространству имён Array API, что и обучающие данные. Например, если для обучения использовалось пространство имён CUDA в PyTorch, атрибуты обученной модели будут находиться на GPU. Передача данных из другого пространства имён или с другого устройства в том же пространстве имён в transform или predict приводит к ошибке:

>>> with config_context(array_api_dispatch=True):
...     lda.transform(X_np)
Traceback (most recent call last):
    ...
ValueError: Inputs passed to LinearDiscriminantAnalysis.transform() must use the same namespace and the same device as those passed to fit()...

>>> from sklearn.utils._array_api import move_estimator_to
>>> import numpy as np
>>> lda_np = move_estimator_to(lda, np, device="cpu")
>>> with config_context(array_api_dispatch=True):
...     X_trans = lda_np.transform(X_np)
>>> type(X_trans)
<class 'numpy.ndarray'>

13.2.4. Поддержка входных данных, совместимых с Array API

Оцениватели и другие инструменты scikit-learn, поддерживающие входные данные, совместимые с Array API.

13.2.5. Обработка типов входных и выходных массивов

Оцениватели и функции оценки могут принимать входные массивы из разных библиотек массивов и/или с разных устройств. Если передан набор входных массивов разных типов, scikit-learn при необходимости преобразует массивы, чтобы привести их к единому виду.

Для оценивателей действует правило «всё следует за X » — входные массивы разных типов преобразуются так, чтобы соответствовать библиотеке массивов и устройству X. Для функций оценки действует правило «всё следует за y_pred » — входные массивы разных типов преобразуются так, чтобы соответствовать библиотеке массивов и устройству y_pred.

Если функция или метод вызывается с входными данными, совместимыми с Array API, принято возвращать массивы из той же библиотеки массивов и на том же устройстве, что и входные данные.

13.2.6. Общие проверки оценивателей (для разработчиков)

Добавьте тег array_api_support в набор тегов оценивателя, чтобы указать, что он поддерживает Array API. Это включит специальные проверки в составе общих тестов, позволяющие убедиться, что результаты оценивателей одинаковы при использовании стандартного NumPy и входных данных Array API.

Для запуска этих проверок необходимо установить array-api-strict в тестовом окружении. Это позволяет запускать проверки без GPU. Для запуска проверок на настоящих GPU-устройствах также необходимо установить PyTorch, CuPy и/или dpnp, а также иметь совместимое оборудование GPU. Ожидается, что полный охват GPU будет распределён между несколькими машинами, поскольку для бэкендов CUDA, MPS и Intel GPU требуется разное оборудование. Проверки, которые невозможно выполнить или для которых отсутствуют зависимости, будут автоматически пропущены. Поэтому важно запускать тесты с флагом -v, чтобы видеть, какие проверки пропускаются:

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

Запуск тестов scikit-learn на array-api-strict должен помочь выявить большинство проблем в коде, связанных с обработкой входных данных с нескольких устройств, за счёт использования имитируемых устройств, отличных от CPU. Это позволяет быстро выполнять итеративную разработку и отладку кода, связанного с Array API.

Однако для полной проверки обработки входных данных PyTorch, CuPy или dpnp, размещённых на реальных GPU-устройствах, необходимо запускать тесты с этими библиотеками и оборудованием. Это можно сделать с помощью Google Colab для CUDA или используя нашу инфраструктуру CI для pull request. Тесты для CUDA и Intel GPU запускаются сопровождающими вручную. Тестирование Intel GPU для PyTorch XPU и dpnp выполняется на выделенном собственном runner: probabl-ai/scikit-learn-intel-workflow.

13.2.7. Примечания

PYTORCH_ENABLE_MPS_FALLBACK=1 pytest -k "array_api" -v