Skip to content

API reference: XGBoost integration#

You can use the Neptune integration with XGBoost to capture model training metadata with NeptuneCallback.


For an in-depth tutorial, see IntegrationsWorking with XGBoost


Neptune callback for logging metadata during XGBoost model training.


This callback requires xgboost>=1.3.0.

The callback logs the following:

  • Metrics
  • All parameters
  • Learning rate
  • The pickled model
  • Visualizations
  • If early stopping is activated, best_score and best_iteration are also logged.

Metrics are logged for every dataset in the evals list and for every metric specified.

Example: With evals = [(dtrain, "train"), (dval, "valid")] and "eval_metric": ["mae", "rmse"], four metrics are created: "train/mae", "train/rmse", "valid/mae", and "valid/rmse".

The callback works with the xgboost.train() and functions, and with the scikit-learn API


Name       Type Default     Description
run Run - An existing run reference, as returned by neptune.init_run().
base_namespace str, optional "training" Namespace under which all metadata logged by the Neptune callback will be stored.
log_model bool True Whether to log the model as a pickled file at the end of training.
log_importance bool True Whether to log feature importance charts at the end of training.
max_num_features int 10 Max number of top features to log on the importance charts. Works when log_importances is set to True. If None or <1, all features will be displayed.

For details, see xgboost.plot_importance() .

log_tree list of int None Indexes of target trees to log as charts. Requires the Graphviz library to be installed.

For details, see xgboost.to_graphviz() .

tree_figsize int 30 Controls the size of the visualized tree image. Increase this in case you work with large trees. Works when log_trees is not None.


# Create a run
import as neptune
run = neptune.init_run(project="workspace-name/project-name")  # (1)

# Create Neptune callback and pass it to xgb.train()
from import NeptuneCallback
neptune_callback = NeptuneCallback(run=run)

xgb.train( ..., callbacks=[neptune_callback])

# When creating the callback, you can specify what you want to log and where
neptune_callback = NeptuneCallback(
    log_tree=[0, 1, 2, 3],
  1. The full project name. For example, "ml-team/classification". To copy it, navigate to the project settingsProperties.