Aller au contenu principal

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

Ouvrir dans Databricks

Entraînez un modèle de régression XGBoost sur un seul GPU à l'aide de AI Runtime. L'accélération GPU accélère considérablement l'entraînement de modèles par rapport à l'entraînement basé sur le CPU, en particulier pour les grands datasets.

Concepts clés abordés :

  • Entraînement accéléré par GPU : Utilise la méthode d'arbre hist de XGBoost avec un périphérique CUDA pour un entraînement plus rapide.
  • Point de contrôle du modèle : enregistre l'état du modèle périodiquement dans les volumes Unity Catalog pour la récupération et la formation incrémentielle
  • Dataset California Housing : Une tâche de régression prévoyant les valeurs médianes des logements

Pour plus d'information, consultez Prise en charge du GPU XGBoost et volumes Unity Catalog.

Se connecter au compute GPU serverless

Cliquez sur le menu déroulant Connecter et sélectionnez Serverless GPU . Ouvrez le panneau latéral Environnement , définissez Accélérateur sur 1xA10 et sélectionnez AI v5 .

Configurer l'emplacement du point de contrôle Unity Catalog

Définissez l’emplacement du volume Unity Catalog où les points de contrôle du modèle seront enregistrés. Le Notebook utilise des parameters query pour configurer le catalogue, le schéma, le volume et le nom du modèle.

Python
# 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 une fonction de rappel de checkpoint qui enregistre l'état du modèle toutes les 50 manches de boosting dans le volume Unity Catalog. Cela permet la récupération après les pannes et l'entraînement incrémental.

Python
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,
interval=50, # save every 50 boosting rounds
)

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

Ouvrir dans Databricks

Chargez le dataset California Housing, configurez XGBoost pour l'entraînement sur GPU et entraînez un modèle de régression. Le modèle prédit les valeurs médianes des maisons en utilisant des caractéristiques telles que l'emplacement, le nombre de pièces et la densité de population.

Python
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é du 150e tour de boosting et évaluez ses performances. Ceci montre comment reprendre l'entraînement ou utiliser les états intermédiaires du modèle.

Python
# 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