Melatih model XGBoost pada satu GPU

Buka versi buku catatan halaman ini

Notebook ini menunjukkan cara melatih model regresi XGBoost pada satu GPU menggunakan komputasi GPU tanpa server Databricks. Akselerasi GPU secara signifikan mempercepat pelatihan model dibandingkan dengan pelatihan berbasis CPU, terutama untuk himpunan data besar.

Konsep utama tercakup:

• Pelatihan yang dipercepat GPU: Menggunakan metode pohon XGBoost hist dengan perangkat CUDA untuk pelatihan yang lebih cepat
• Model checkpointing: Menyimpan status model secara berkala ke volume Katalog Unity untuk pemulihan dan pelatihan bertahap
• Himpunan data Perumahan California: Tugas regresi yang memprediksi nilai rumah median

Untuk informasi selengkapnya, lihat Dukungan GPU XGBoost dan Volume Katalog Unity.

Persyaratan

Notebook ini memerlukan kluster komputasi berkemampuan GPU. Komputasi GPU tanpa server Databricks dipilih secara otomatis saat menjalankan sel.

Menginstal pustaka yang diperlukan

Instal XGBoost versi 2.0.3 dan scikit-learn untuk metrik pemuatan dan evaluasi himpunan data.

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

Verifikasi bahwa XGBoost 2.0.3 diinstal dengan benar.

%pip show xgboost

Mengonfigurasi lokasi titik pemeriksaan Katalog Unity

Tentukan lokasi volume Katalog Unity tempat titik pemeriksaan model akan disimpan. Notebook menggunakan parameter kueri untuk mengonfigurasi katalog, skema, volume, dan nama model.

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

Buat panggilan balik titik pemeriksaan yang menyimpan status model setiap 50 putaran peningkatan ke volume Katalog Unity. Ini memungkinkan pemulihan dari kegagalan dan pelatihan inkremental.

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
)

Melatih model XGBoost pada satu GPU

Buka versi buku catatan halaman ini

Muat himpunan data California Housing, konfigurasikan XGBoost untuk pelatihan GPU, dan latih model regresi. Model ini memprediksi nilai rumah median menggunakan fitur seperti lokasi, jumlah kamar, dan kepadatan populasi.

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

Muat model dari titik pemeriksaan dan evaluasi

Muat titik pemeriksaan yang disimpan sebelumnya dari putaran peningkatan ke-150 dan evaluasi performanya. Ini menunjukkan cara melanjutkan pelatihan atau menggunakan status model perantara.

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

Langkah berikutnya

• Dokumentasi Dukungan GPU XGBoost
• Praktik terbaik untuk komputasi GPU Tanpa Server
• Memecahkan masalah pada komputasi GPU tanpa server
• Pelatihan terdistribusi multi-GPU dan multi-simpul
• Volume pada Unity Catalog

Contoh buku catatan

Melatih model XGBoost pada satu GPU

Dapatkan buku catatan