Partage via

Entraîner un modèle XGBoost sur un seul GPU

Ouvrir la version du bloc-notes de cette page

Ce notebook montre comment entraîner un modèle de régression XGBoost sur un seul GPU en utilisant la puissance de calcul GPU serverless de Databricks. L’accélération GPU accélère considérablement l’entraînement du modèle par rapport à l’entraînement basé sur le processeur, en particulier pour les jeux de données volumineux.

Concepts clés abordés :

  • Formation accélérée par GPU : utilise la méthode d’arborescence de hist XGBoost avec l’appareil CUDA pour une formation plus rapide
  • Enregistrement de point de contrôle du modèle : enregistre régulièrement l’état du modèle dans les volumes du catalogue Unity pour la récupération et l’apprentissage incrémentiel
  • Jeu de données de logement californien : problème de régression pour la prévision des valeurs médianes des maisons

Pour plus d’informations, consultez Prise en Charge du GPU XGBoost et Volumes du Catalogue Unity.

Exigences

Ce notebook nécessite un cluster de calcul avec GPU. Le calcul GPU serverless Databricks est automatiquement sélectionné lors de l’exécution de cellules.

Installer les bibliothèques requises

Installez XGBoost version 2.0.3 et scikit-learn pour le chargement et les métriques d’évaluation du jeu de données.

%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()

Vérifiez que XGBoost 2.0.3 est installé correctement.

%pip show xgboost

Configurer l’emplacement du point de contrôle du catalogue Unity

Définissez l’emplacement du volume du catalogue Unity où les points de contrôle de modèle seront enregistrés. Le notebook utilise des paramètres de requête pour configurer le catalogue, le schéma, le volume et le nom du modèle.

# 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}")

Créez un rappel de point de contrôle qui enregistre l’état du modèle toutes les 50 rondes d’amélioration dans le volume du catalogue Unity. Cela permet la récupération à partir d’échecs et de l’entraînement incrémentiel.

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
)

Entraîner un modèle XGBoost sur un seul GPU

Ouvrir la version du bloc-notes de cette page

Chargez le jeu de données California Housing, configurez XGBoost pour l’entraînement GPU et entraînez un modèle de régression. Le modèle prédit les valeurs médianes des maisons à l’aide de caractéristiques telles que l’emplacement, le nombre de pièces et la densité de population.

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}")

Charger le modèle à partir du point de contrôle et évaluer

Chargez un point de contrôle précédemment enregistré à partir de la 150e itération de boosting et évaluer ses performances. Cela montre comment reprendre l’entraînement ou utiliser des états de modèle intermédiaires.

# 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}")

Étapes suivantes

Exemple de notebook

Entraîner un modèle XGBoost sur un seul GPU

Obtenir un ordinateur portable

  • Last updated on