Trénování modelu XGBoost na jednom GPU

Otevřít verzi poznámkového bloku této stránky

Trénování regresního modelu XGBoost na jednom GPU pomocí modulu runtime AI Akcelerace GPU výrazně zrychluje trénování modelu v porovnání s trénováním založeným na procesoru, zejména u velkých datových sad.

Probírané klíčové koncepty:

• Trénování urychlené pomocí GPU: Pro rychlejší trénování používá XGBoostovu metodu stromů s využitím zařízení CUDA.
• Vytváření kontrolních bodů modelu: Pravidelně ukládá stav modelu do svazků katalogu Unity pro obnovení a přírůstkové trénování.
• Kalifornský dataset pro bydlení: Regresní úloha předpovídající střední hodnoty cen domů

Další informace najdete v podpoře GPU XGBoost a Unity Catalog - svazky.

Připojení k bezserverovému výpočetnímu prostředí GPU

Klikněte na rozevírací seznam Připojit a vyberte Bezserverové GPU. Otevřete boční panel Prostředí , nastavte Akcelerátor na 1xA10 a vyberte AI v5.

Konfigurace umístění kontrolního bodu katalogu Unity

Definujte umístění svazku katalogu Unity, kam se budou ukládat kontrolní body modelu. Notebook používá parametry dotazu (query) ke konfiguraci katalogu, schématu, objemu a názvu modelu.

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

Vytvořte zpětné volání pro kontrolní bod, které uloží stav modelu každých 50 kol zvyšování výkonu na svazek katalogu Unity. To umožňuje zotavení ze selhání a přírůstkového trénování.

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,
    interval=50,         # save every 50 boosting rounds
)

Trénování modelu XGBoost na jednom GPU

Otevřít verzi poznámkového bloku této stránky

Načtěte datovou sadu o bydlení v Kalifornii, nakonfigurujte XGBoost pro využití GPU při trénování a natrénujte regresní model. Model predikuje medián hodnot domu pomocí funkcí, jako je umístění, počet místností a hustota populace.

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

Načíst model z kontrolního bodu a provést vyhodnocení

Načtěte dříve uložený kontrolní bod z 150. kola boosting a vyhodnoťte jeho výkon. To ukazuje, jak pokračovat v trénování nebo používat stavy zprostředkujícího modelu.

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

Další kroky

• Dokumentace podpory GPU XGBoost
• Osvědčené postupy pro výpočetní prostředí GPU bez serveru
• Řešení potíží s výpočetními prostředky GPU bez serveru
• Distribuované trénování s více GPU a více uzly
• Svazky katalogu Unity

Ukázkový poznámkový blok

Trénování modelu XGBoost na jednom GPU

Pořiďte si notebook