Regola la classificazione XGBoost su GPU con Optuna

Questo notebook dimostra l'ottimizzazione degli iperparametri per un modello di classificazione XGBoost su una singola GPU utilizzando Optuna e Databricks AI Runtime.

Connettersi alla GPU serverless

1. Clicca sul menu a tendina Collega nella barra degli strumenti del taccuino.
2. Selezionare GPU serverless.
3. Apri il pannello laterale Environment.
4. Imposta Acceleratore su 1xA10.
5. Seleziona l'AI di base dell'ambiente v6 o superiore.

Configura il tracciamento MLflow

Seguendo il modello di esempio end-to-end di Databricks MLflow, questo notebook traccia le esecuzioni di tuning Optuna e il modello finale XGBoost con MLflow. Ogni prova di tuning viene registrata come una run annidata, e il modello finale è registrato come un artefatto MLflow deployabile.

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

Caricare e preparare i dati

Utilizziamo il dataset Breast Cancer Wisconsin di scikit-learn, un compito di classificazione binaria con 30 caratteristiche numeriche e 569 campioni.

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

Definisci la funzione obiettivo di Optuna

Optuna effettua la ricerca sugli iperparametri XGBoost utilizzando l'addestramento accelerato tramite GPU (device: "cuda"). Seguendo la struttura di esempio MLflow di Databricks, ogni tentativo viene anche catturato come una run MLflow annidata, e la potatura interrompe prematuramente i tentativi poco promettenti.

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

Esegui l'ottimizzazione degli iperparametri

Creiamo uno studio Optuna per minimizzare la perdita di logaritmi di validazione e eseguiamo 50 prove. Una run MLflow genitore traccia la sessione di tuning, mentre ogni trial viene registrata come una run figlia annidata.

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

Modello finale di addestramento con i migliori parametri

Addestrare il modello finale XGBoost usando i migliori iperparametri trovati da Optuna, valutarlo sul set di test offerto e registrare un artefatto MLflow deployabile con firma ed esempio di input.

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

Visualizza i risultati dell'ottimizzazione

Optuna fornisce strumenti di visualizzazione integrati per comprendere il processo di ottimizzazione e l'importanza degli iperparametri. Questi completano le esecuzioni MLflow registrate durante la messa a punto e l'addestramento finale del modello.

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

Passaggi successivi

• Esempio end-to-end di Databricks MLflow
• Documentazione della GPU XGBoost
• Documentazione Optuna
• Databricks Calcolo GPU serverless
• Modello di registro nel Catalogo Unity

Notebook di esempio

Regola la classificazione XGBoost su GPU con Optuna

Ottieni il notebook