Visualizing XGBoost Trees

XGBoost is a popular gradient-boosting library for building regression and classification models. The core of XGBoost is an ensemble of decision trees. In this post, we'll look at how to visualize and interpret individual trees from an XGBoost model.

Let's start by loading a simple sample dataset from sci-kit-learn - the diabetes dataset. This dataset has 13 features for predicting house prices.


import xgboost as xgb
from sklearn.datasets import load_diabetes
from sklearn.model_selection import train_test_split
import matplotlib.pyplot as plt

# Load diabetes dataset
X, y = load_diabetes(return_X_y=True)

Load dataset

# Train XGBoost model
model = xgb.XGBRegressor()
model.fit(X_train, y_train)

Train a model


# Visualize first tree
xgb.plot_tree(model, num_trees=0)
plt.rcParams['figure.figsize'] = [100, 100]
plt.show()

Visualize the trained model