Обучение модели XGBoost на одном GPU

Откройте версию страницы для записной книжки

В этой записной книжке показано, как обучить модель регрессии XGBoost на одном GPU с помощью бессерверных вычислений GPU Databricks. Ускорение GPU значительно ускоряет обучение модели по сравнению с обучением на основе ЦП, особенно для больших наборов данных.

Основные понятия:

• Обучение с ускорением GPU: использует метод дерева XGBoost hist с устройством CUDA для ускорения обучения
• Контрольная точка модели: периодически сохраняет состояние модели в томах каталога Unity для восстановления и добавочного обучения
• Набор данных о жилье в Калифорнии: задача регрессии по прогнозированию медианных значений стоимости домов

Дополнительные сведения см. в разделе Поддержка GPU XGBoost и Тома каталога Unity.

Требования

Для этой записной книжки требуется вычислительный кластер с поддержкой GPU. При выполнении ячеек в Databricks автоматически выбираются бессерверные вычислительные ресурсы GPU.

Установка необходимых библиотек

Установите XGBoost версии 2.0.3 и scikit-learn для загрузки и оценки набора данных.

%pip install xgboost==2.0.3 # due to this issue: https://github.com/ray-project/xgboost_ray/issues/312
%pip install scikit-learn
dbutils.library.restartPython()

Убедитесь, что XGBoost 2.0.3 установлен правильно.

%pip show xgboost

Настройка расположения контрольной точки каталога Unity

Определите расположение тома каталога Unity, в котором будут сохранены контрольные точки модели. Записная книжка использует параметры запроса для настройки каталога, схемы, тома и имени модели.

# You must have `USE CATALOG` privileges on the catalog, and you must have `USE SCHEMA` privileges on the schema.
# If necessary, change the catalog and schema name here.
dbutils.widgets.text("uc_catalog", "main")
dbutils.widgets.text("uc_schema", "default")
dbutils.widgets.text("uc_model_name", "custom_transformer")
dbutils.widgets.text("uc_volume", "checkpoints")

UC_CATALOG = dbutils.widgets.get("uc_catalog")
UC_SCHEMA = dbutils.widgets.get("uc_schema")
UC_VOLUME = dbutils.widgets.get("uc_volume")
MODEL_NAME = dbutils.widgets.get("uc_model_name")
CHECKPOINT_PATH = f"/Volumes/{UC_CATALOG}/{UC_SCHEMA}/{UC_VOLUME}/{MODEL_NAME}"
CHECKPOINT_PREFIX = "checkpoint"

print(f"UC_CATALOG: {UC_CATALOG}")
print(f"UC_SCHEMA: {UC_SCHEMA}")
print(f"UC_VOLUME: {UC_VOLUME}")
print(f"CHECKPOINT_PATH: {CHECKPOINT_PATH}")

Создайте обратный вызов контрольной точки, который сохраняет состояние модели каждые 50 округлений до тома каталога Unity. Это обеспечивает восстановление после сбоев и добавочного обучения.

import os
from xgboost.callback import TrainingCheckPoint

# Create the UC Volume where the checkpoint will be saved if it doesn't exist already
os.makedirs(CHECKPOINT_PATH, exist_ok=True)

# Create a callback to checkpoint to a UC volume
checkpoint_cb = TrainingCheckPoint(
    directory=CHECKPOINT_PATH,
    name=CHECKPOINT_PREFIX,
    iterations=50,       # save every 50 boosting rounds
)

Обучение модели XGBoost на одном GPU

Откройте версию страницы для записной книжки

Загрузите набор данных о жилье в Калифорнии, настройте XGBoost для обучения на GPU и обучите регрессионную модель. Модель прогнозирует значения медиана дома с помощью таких функций, как расположение, количество комнат и плотность населения.

import xgboost as xgb
from sklearn.datasets import fetch_california_housing
from sklearn.model_selection import train_test_split
from sklearn.metrics import root_mean_squared_error

# Load California Housing dataset
X, y = fetch_california_housing(return_X_y=True)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# Convert to DMatrix
dtrain = xgb.DMatrix(X_train, label=y_train)
dtest = xgb.DMatrix(X_test, label=y_test)

# GPU training parameters for regression
params = {
    "tree_method": "hist",        # Use GPU histogram
    "device": "cuda",
    "objective": "reg:squarederror",  # Regression objective
    "eval_metric": "rmse",            # Root Mean Squared Error
    "max_depth": 6,
    "learning_rate": 0.1,
}

# Train the model
bst = xgb.train(
    params=params,
    dtrain=dtrain,
    num_boost_round=200,
    evals=[(dtest, "eval"), (dtrain, "train")],
    verbose_eval=10,
    callbacks=[checkpoint_cb]
)

# Predict
y_pred = bst.predict(dtest)

# Evaluate
rmse = root_mean_squared_error(y_test, y_pred)
print(f"✅ RMSE on test set: {rmse:.4f}")

Загрузка модели из контрольной точки и оценка

Загрузите ранее сохраненную контрольную точку из 150-го раунда повышения и оцените ее производительность. В этом руководстве показано, как возобновить обучение или использовать промежуточные состояния модели.

# Take sample checkpoint from 150th step
checkpoint = f"{CHECKPOINT_PATH}/{CHECKPOINT_PREFIX}_150.json"

# Load the model from a checkpoint
bst = xgb.Booster()
bst.load_model(checkpoint)

dtest = xgb.DMatrix(X_test)
y_pred = bst.predict(dtest)

# Evaluate
rmse = root_mean_squared_error(y_test, y_pred)
print(f"✅ RMSE on test set: {rmse:.4f}")

Дальнейшие действия

• Документация по поддержке GPU XGBoost
• Рекомендации по вычислению бессерверного GPU
• Устранение неполадок с бессерверными вычислительными процессорами GPU
• Распределенное обучение с несколькими GPU и несколькими узлами
• Тома каталога Unity

Пример записной книжки

Обучение модели XGBoost на одном GPU

Получите ноутбук