Ajustez la classification XGBoost sur GPU avec Optuna

Ce notebook démontre l’optimisation des hyperparamètres pour un modèle de classification XGBoost sur un seul GPU utilisant Optuna et Databricks AI Runtime.

Se connecter à un GPU sans serveur

1. Cliquez sur le menu déroulant Connecter dans la barre d’outils du carnet.
2. Sélectionnez GPU sans serveur.
3. Ouvre le panneau latéral Environnement.
4. Réglez l’accélérateur sur 1xA10.
5. Sélectionnez l’environnement de base AI v6 ou version ultérieure.

Configurer le suivi MLflow

En suivant le modèle MLflow de bout en bout de Databricks, ce notebook suit les exécutions de réglage Optuna et le modèle final XGBoost avec MLflow. Chaque essai de réglage est enregistré comme une exécution imbriquée, et le modèle final est enregistré comme un artefact MLflow déployable.

import mlflow
from mlflow.models import infer_signature
import xgboost as xgb

mlflow.xgboost.autolog(log_models=False)
print("MLflow XGBoost autologging enabled for parameters and metrics.")

Charger et préparer les données

Nous utilisons le jeu de données Breast Cancer Wisconsin de scikit-learn, une tâche de classification binaire avec 30 caractéristiques numériques et 569 échantillons.

import numpy as np
import pandas as pd
from sklearn.datasets import load_breast_cancer
from sklearn.model_selection import train_test_split
import xgboost as xgb

# Load dataset as a DataFrame so the logged MLflow model keeps feature names
data = load_breast_cancer(as_frame=True)
X = data.data
y = data.target

# Hold out a test set used only for the final evaluation, then split the remainder
# into train and validation sets. Tuning uses the validation set, so hyperparameters
# are never selected on the data used to report the final metrics.
X_trainval_df, X_test_df, y_trainval, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)
X_train_df, X_val_df, y_train, y_val = train_test_split(
    X_trainval_df, y_trainval, test_size=0.25, random_state=42, stratify=y_trainval
)

# Convert to XGBoost DMatrix format for efficient training
dtrain = xgb.DMatrix(X_train_df, label=y_train)
dval = xgb.DMatrix(X_val_df, label=y_val)
dtest = xgb.DMatrix(X_test_df, label=y_test)

print(f"Train: {X_train_df.shape[0]}, Validation: {X_val_df.shape[0]}, Test: {X_test_df.shape[0]}")
print(f"Features: {X_train_df.shape[1]}")
print(f"Classes: {np.unique(y)}")
display(X_train_df.head())

Définissons la fonction objectif d’Optuna

Optuna explore les hyperparamètres XGBoost en utilisant l’entraînement accéléré par GPU (device: "cuda"). En suivant la structure d’exemple MLflow de Databricks, chaque essai est également capturé comme une exécution MLflow imbriquée, et l’élagage arrête tôt les essais peu prometteurs.

import optuna
from sklearn.metrics import log_loss
from optuna_integration import XGBoostPruningCallback

def objective(trial):
    params = {
        "tree_method": "hist",
        "device": "cuda",
        "objective": "binary:logistic",
        "eval_metric": "logloss",
        "max_depth": trial.suggest_int("max_depth", 3, 10),
        "learning_rate": trial.suggest_float("learning_rate", 1e-3, 0.3, log=True),
        "subsample": trial.suggest_float("subsample", 0.5, 1.0),
        "colsample_bytree": trial.suggest_float("colsample_bytree", 0.5, 1.0),
        "min_child_weight": trial.suggest_int("min_child_weight", 1, 10),
        "gamma": trial.suggest_float("gamma", 1e-8, 1.0, log=True),
        "reg_alpha": trial.suggest_float("reg_alpha", 1e-8, 10.0, log=True),
        "reg_lambda": trial.suggest_float("reg_lambda", 1e-8, 10.0, log=True),
    }

    n_estimators = trial.suggest_int("n_estimators", 50, 500)

    with mlflow.start_run(nested=True, run_name=f"optuna-trial-{trial.number}"):
        mlflow.set_tags({
            "phase": "tuning",
            "trial_number": trial.number,
            "model_family": "xgboost",
            "device": "cuda",
        })
        mlflow.log_param("n_estimators", n_estimators)

        # Evaluate and prune on the validation set; the test set is reserved for the final evaluation.
        pruning_callback = XGBoostPruningCallback(trial, "validation-logloss")
        model = xgb.train(
            params,
            dtrain,
            num_boost_round=n_estimators,
            evals=[(dval, "validation")],
            callbacks=[pruning_callback],
            verbose_eval=False,
        )

        preds = model.predict(dval)
        trial_logloss = log_loss(y_val, preds)
        mlflow.log_metric("trial_logloss", trial_logloss)
        return trial_logloss

Exécuter l'optimisation des hyperparamètres

Nous créons une étude Optuna pour minimiser la perte de journalisation de validation et réalisons 50 essais. Une exécution MLflow parent suit la session de réglage, tandis que chaque essai est enregistré comme une exécution enfant imbriquée.

optuna.logging.set_verbosity(optuna.logging.WARNING)

with mlflow.start_run(run_name="optuna-xgboost-gpu-classification") as tuning_run:
    mlflow.set_tags({
        "phase": "hyperparameter_tuning",
        "model_family": "xgboost",
        "task": "binary_classification",
        "device": "cuda",
        "optimizer": "optuna",
    })
    mlflow.log_param("n_trials", 50)

    study = optuna.create_study(direction="minimize", study_name="xgboost-gpu-tuning")
    study.optimize(objective, n_trials=50, show_progress_bar=True)

    mlflow.log_metric("best_trial_logloss", study.best_trial.value)
    mlflow.log_params({f"best_{key}": value for key, value in study.best_trial.params.items()})
    tuning_run_id = tuning_run.info.run_id

print(f"\nBest trial logloss: {study.best_trial.value:.6f}")
print("Best hyperparameters:")
for key, value in study.best_trial.params.items():
    print(f"  {key}: {value}")
print(f"\nMLflow tuning run_id: {tuning_run_id}")

Modèle final de train avec les meilleurs paramètres

Entraîner le modèle XGBoost final en utilisant les meilleurs hyperparamètres trouvés par Optuna, l’évaluer sur l’ensemble de test posé, puis enregistrer un artefact MLflow déployable avec signature et exemple d’entrée.

from sklearn.metrics import accuracy_score, classification_report, confusion_matrix, log_loss
import matplotlib.pyplot as plt

# Build final parameters from the best trial
best_params = study.best_trial.params.copy()
n_estimators = best_params.pop("n_estimators")

best_params.update({
    "tree_method": "hist",
    "device": "cuda",
    "objective": "binary:logistic",
    "eval_metric": "logloss",
})

# Refit on train + validation (all non-test data) with the best hyperparameters,
# then evaluate once on the held-out test set.
X_fit_df = pd.concat([X_train_df, X_val_df])
y_fit = pd.concat([y_train, y_val])
dfit = xgb.DMatrix(X_fit_df, label=y_fit)

with mlflow.start_run(run_name="best-xgboost-gpu-model") as final_run:
    mlflow.set_tags({
        "phase": "final_model",
        "model_family": "xgboost",
        "task": "binary_classification",
        "device": "cuda",
    })
    mlflow.log_param("n_estimators", n_estimators)
    mlflow.log_params(best_params)

    # Train the final model on all non-test data
    final_model = xgb.train(
        best_params,
        dfit,
        num_boost_round=n_estimators,
        evals=[(dtest, "test")],
        verbose_eval=False,
    )

    # Predict and evaluate on the held-out test set
    y_pred_proba = final_model.predict(dtest)
    y_pred = (y_pred_proba > 0.5).astype(int)

    accuracy = accuracy_score(y_test, y_pred)
    test_logloss = log_loss(y_test, y_pred_proba)
    report_text = classification_report(y_test, y_pred, target_names=data.target_names)
    report_dict = classification_report(
        y_test, y_pred, target_names=data.target_names, output_dict=True
    )
    conf_matrix = confusion_matrix(y_test, y_pred)

    mlflow.log_metric("test_accuracy", accuracy)
    mlflow.log_metric("test_logloss", test_logloss)
    mlflow.log_dict(report_dict, "classification_report.json")

    fig, ax = plt.subplots(figsize=(4, 4))
    image = ax.imshow(conf_matrix, cmap="Blues")
    ax.set_title("Confusion Matrix")
    ax.set_xlabel("Predicted label")
    ax.set_ylabel("True label")
    ax.set_xticks([0, 1])
    ax.set_yticks([0, 1])
    ax.set_xticklabels(data.target_names)
    ax.set_yticklabels(data.target_names)
    for row_idx in range(conf_matrix.shape[0]):
        for col_idx in range(conf_matrix.shape[1]):
            ax.text(col_idx, row_idx, conf_matrix[row_idx, col_idx], ha="center", va="center")
    fig.colorbar(image, ax=ax)
    plt.tight_layout()
    mlflow.log_figure(fig, "confusion_matrix.png")
    plt.close(fig)

    signature = infer_signature(X_test_df, y_pred_proba)
    model_info = mlflow.xgboost.log_model(
        final_model,
        name="model",
        signature=signature,
        input_example=X_fit_df.head(3),
    )
    final_run_id = final_run.info.run_id

print(f"Test Accuracy: {accuracy:.4f}")
print(f"Test Log Loss: {test_logloss:.6f}\n")
print("Classification Report:")
print(report_text)
print("\nConfusion Matrix:")
print(conf_matrix)
print(f"\nFinal MLflow run_id: {final_run_id}")
print(f"Model URI: {model_info.model_uri}")

Visualiser les résultats d’optimisation

Optuna propose des outils de visualisation intégrés pour comprendre le processus d’optimisation et l’importance des hyperparamètres. Ces éléments complètent les exécutions MLflow enregistrées lors de l’ajustement et de l’entraînement final du modèle.

import matplotlib.pyplot as plt
from optuna.visualization.matplotlib import plot_optimization_history, plot_param_importances

plt.figure(figsize=(10, 4))
plot_optimization_history(study, target_name="Log Loss")
plt.tight_layout()
plt.show()

plt.figure(figsize=(10, 4))
plot_param_importances(study)
plt.tight_layout()
plt.show()

Étapes suivantes

• Exemple de bout en bout MLflow sur Databricks
• Documentation sur les GPU XGBoost
• Documentation Optuna
• Calcul GPU Databricks Serverless
• Modèle de registre dans le catalogue Unity

Exemple de bloc-notes

Ajustez la classification XGBoost sur GPU avec Optuna

Obtenir un ordinateur portable