Ajustez la classification XGBoost sur GPU avec Optuna

Ouvrir dans Databricks

Ce notebook présente l'optimisation des hyperparamètres pour un modèle de classification XGBoost sur un seul GPU à l'aide de Optuna et de Databricks AI Runtime.

remarque

Cet exemple requiert la version 6 ou ultérieure de l'environnement Databricks AI.

Se connecter au GPU serverless

1. Cliquez sur le menu déroulant Connect dans la barre d'outils du notebook.
2. Sélectionnez Serverless GPU .
3. Ouvrez le panneau latéral Environnement .
4. Définissez Accelerator sur 1xA10.
5. Sélectionnez l'environnement de base AI v6 ou une version ultérieure.

Configurer le suivi MLflow

Conformément au modèle d’exemple de bout en bout de Databricks avec MLflow, ce notebook suit les exécutions d’ajustement d’Optuna et le modèle XGBoost final avec MLflow. Chaque essai d’ajustement est enregistré en tant qu’exécution imbriquée, et le modèle final est enregistré en tant qu’artefact de modèle MLflow déployable.

Python

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

Chargez et préparez les données

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

Python

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éfinissez la fonction objectif d'Optuna

Optuna explore les hyperparamètres XGBoost à l’aide d’un entraînement accéléré par GPU (device: "cuda"). Selon la structure d’exemple Databricks avec MLflow, chaque essai est également capturé comme une exécution MLflow imbriquée, et l’élagage interrompt prématurément les essais peu prometteurs.

Python

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écutez l'optimisation des hyperparamètres

Nous créons une étude Optuna pour minimiser la perte logarithmique de validation et exécutons 50 essais. Une exécution MLflow parente suit la session d’ajustement, tandis que chaque essai est enregistré en tant qu’exécution enfant imbriquée.

Python

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

Entraîner le modèle final avec les meilleurs paramètres

Entraînez le modèle XGBoost final à l'aide des meilleurs hyperparamètres trouvés par Optuna, évaluez-le sur le jeu de test isolé, et enregistrez un artefact de modèle MLflow déployable avec la signature et l'exemple d'entrée.

Python

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

Visualisez les résultats de l'optimisation

Optuna fournit des outils de visualisation intégrés pour comprendre le processus d'optimisation et l'importance des hyperparamètres. Ceux-ci complètent les run MLflow enregistrés lors de l'ajustement et de l'entraînement final du modèle.

Python

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 de Databricks MLflow
• Documentation XGBoost GPU
• Documentation d'Optuna
• Databricks Serverless GPU compute
• Enregistrer le modèle dans Unity Catalog

Exemple de notebook

Ajuster la classification XGBoost sur GPU avec Optuna

Ouvrir le notebook dans un nouvel ongletOuvrir dans Databricks