XGBoost

Integrate W&B with XGBoost to log gradient boosting metrics, feature importance, and model performance automatically.

This page shows you how to use the W&B integration with XGBoost to automatically log gradient boosting metrics, model configurations, feature importance, and trained boosters so you can track, compare, and reproduce your XGBoost experiments.

The W&B Python SDK (wandb) has a WandbCallback callback that logs metrics, configs, and saved boosters from training with XGBoost. See a live W&B Dashboard with outputs from the XGBoost WandbCallback.

Get started

To log XGBoost metrics, configs, and booster models to W&B, pass the WandbCallback to XGBoost:

import wandb
from wandb.integration.xgboost import WandbCallback
from xgboost import XGBClassifier

...
# Start a wandb run
with wandb.init() as run:
  # Pass WandbCallback to the model
  bst = XGBClassifier()
  bst.fit(X_train, y_train, callbacks=[WandbCallback(log_model=True)])

For a comprehensive look at logging with XGBoost and W&B, see the XGBoost and W&B logging notebook.

WandbCallback reference

Functionality

Passing WandbCallback to an XGBoost model does the following:

• Logs the booster model configuration to W&B.
• Logs evaluation metrics collected by XGBoost, such as rmse, accuracy, and so on to W&B.
• Logs training metrics collected by XGBoost (if you provide data to eval_set).
• Logs the best score and the best iteration.
• Saves and uploads your trained model to W&B Artifacts (when log_model = True).
• Logs the feature importance plot when log_feature_importance=True (default).
• Captures the best eval metric in wandb.Run.summary when define_metric=True (default).

Arguments

• log_model: (boolean) if True, saves and uploads the model to W&B Artifacts.

• log_feature_importance: (boolean) if True, logs a feature importance bar plot.

• importance_type: (str) one of {weight, gain, cover, total_gain, total_cover} for tree model. weight for linear model.

• define_metric: (boolean) if True (default), captures model performance at the best step, instead of the last step, of training in your run.summary.

Review the source code for WandbCallback.

For more examples, see the repository of examples on GitHub.

Tune your hyperparameters with Sweeps

W&B Sweeps is a toolkit for configuring, orchestrating, and analyzing hyperparameter testing experiments. This section shows how to combine the XGBoost integration with W&B Sweeps to search across hyperparameter configurations.

To improve model performance, tune hyperparameters like tree depth and learning rate.

You can also try this XGBoost and Sweeps Python script.