単一のGPUでXGBoostモデルをトレーニングする
このノートブックでは、 Databricksサーバレス GPU コンピュートを使用して単一の GPU でXGBoost回帰モデルをトレーニングする方法を示します。 GPUアクセラレーションは、特に大規模データセットの場合、CPUベースのトレーニングと比較して、モデルのトレーニングを大幅に高速化します。
主な概念：
- GPUアクセラレーションによるトレーニング ：CUDAデバイスを使用してXGBoostの
histツリーメソッドでトレーニングを高速化します。 - モデルチェックポイント ：モデルの状態を定期的にUnity Catalogボリュームに保存し、復元と増分更新に備えます。
- カリフォルニア州住宅データセット ：住宅価格の中央値を予測する回帰分析タスク
詳細については、 XGBoost GPU サポートとUnity Catalogボリューム」を参照してください。
要件
このノートブックには、GPU 対応のコンピュート クラスターが必要です。 セルを実行すると、 Databricksサーバレス GPU コンピュートが自動的に選択されます。
必要なライブラリをインストールします
データセットの読み込みとメトリクスの評価のために、 XGBoostバージョン 2.0.3 とscikit-learnをインストールします。
Python
%pip install xgboost==2.0.3 # due to this issue: https://github.com/ray-project/xgboost_ray/issues/312
%pip install scikit-learn
dbutils.library.restartPython()
XGBoost 2.0.3が正しくインストールされていることを確認してください。
Python
%pip show xgboost
Unity Catalogのチェックポイントの位置を設定する
モデルのチェックポイントを保存するUnity Catalogボリュームの場所を定義します。 ノートブックはクエリーを使用して、カタログ、スキーマ、ボリューム、およびモデル名を構成します。
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}")
50回のブーストラウンドごとにモデルの状態をUnity Catalogボリュームに保存するチェックポイントコールバックを作成します。 これにより、失敗からの回復と段階的なトレーニングが可能になります。
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,
iterations=50, # save every 50 boosting rounds
)
単一のGPUでXGBoostモデルをトレーニングする
カリフォルニア住宅データセットをロードし、GPU トレーニング用にXGBoostを構成し、回帰モデルをトレーニングします。 このモデルは、立地、部屋数、人口密度などの特徴量を用いて、住宅価格の中央値を予測します。
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}")
チェックポイントからモデルをロードして評価する
150回目のブーストラウンドで保存したチェックポイントを読み込み、その性能を評価してください。これは、トレーニングを再開する方法、または中間モデル状態を使用する方法を示しています。
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}")
次のステップ
- XGBoost GPU サポートに関するドキュメント
- サーバレス GPU コンピュートのベスト プラクティス
- サーバーレス GPU コンピュートの問題のトラブルシューティング
- マルチGPUおよびマルチノード分散トレーニング
- Unity Catalogボリューム