Optunaを使ったGPUでのXGBoost分類のチューニング

このノートブックでは、Optuna と Databricks AI Runtime を使用して、単一の GPU 上で XGBoost 分類モデルのハイパーパラメータ最適化を行う方法を示しています。

Note

この例はDatabricks AI環境バージョン6以上が必要です。

サーバーレス GPU に接続する

ノートブックのツールバーの「 Connect 」ドロップダウンをクリックしてください。
[サーバーレス GPU] を選択します。
環境サイドパネルを開きます。
アクセラレータを1xA10に設定してください。
ベース環境の AI v6 以上に選択します。

MLflowトラッキングの設定

Databricks MLflow のエンドツーエンドのサンプルパターンに従い、このノートブックでは、MLflow を使用して Optuna のチューニング実行と最終的な XGBoost モデルを追跡します。各チューニングトライアルはネストされた実行としてログされ、最終モデルはデプロイ可能なMLflowモデルのアーティファクトとしてログされます。

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

データの読み込みと準備

ここでは、scikit-learn の Breast Cancer Wisconsin データセットを使用します。これは、30 個の数値特徴量と 569 件のサンプルからなる二値分類タスクです。

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

Optunaの目的関数を定義します

OptunaはGPU加速トレーニング(device: "cuda")を用いてXGBoostのハイパーパラメータを検索します。 Databricks の MLflow サンプル構造に従い、各トライアルもネストされた MLflow 実行として記録され、プルーニングによって見込みの薄いトライアルは早期に打ち切られます。

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

ハイパーパラメータ最適化を実行する

検証ログの損失を最小限に抑えるために Optuna の研究を作成し、50 回の試行を実行します。親MLflowランはチューニングセッションを追跡し、各試行はネストされた子ランとして記録されます。

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

最良パラメータで最終モデルを訓練する

Optunaで見つけた最良のハイパーパラメータを使って最終的なXGBoostモデルを訓練し、保留されたテストセットで評価し、シグネチャと入力例を付けた展開可能なMLflowモデルのアーティファクトをログに記録します。

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

最適化結果を可視化する

Optunaは最適化プロセスやハイパーパラメータの重要性を理解するための組み込みの可視化ツールを提供します。これらはチューニングや最終モデルトレーニング時に記録されたMLflowの実行を補完します。

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

次のステップ

ノートブックの例

Optunaを使ったGPUでのXGBoost分類のチューニング

ノートブックを入手

フィードバック

このページはお役に立ちましたか?

最終更新: 2026-10-08

Optunaを使ったGPUでのXGBoost分類のチューニング

サーバーレス GPU に接続する

MLflowトラッキングの設定

データの読み込みと準備

Optunaの目的関数を定義します

ハイパーパラメータ最適化を実行する

最良パラメータで最終モデルを訓練する

最適化結果を可視化する

次のステップ

ノートブックの例

Optunaを使ったGPUでのXGBoost分類のチューニング

フィードバック

その他のリソース