How to plot the decision trees from XGBoost classifier

In this article, I am going to show you how to plot the decision trees generated by XGBoost models. First, we have to install graphviz (both python library and executable files)

1
2
!pip install graphviz
!apt-get install graphviz

When the graphviz library is installed, we can train an XGBoost model (in this example, I am going to train it using the Titanic dataset).

1
2
3
4
5
6
7
8
from xgboost import XGBClassifier
from xgboost import plot_tree
#(...) other imports

#(...) loading the dataset and data preprocessing

model = XGBClassifier()
model.fit(X_train, y_train, verbose=True, eval_set=[(X_test, y_test)])

To display the trees, we have to use the plot_tree function provided by XGBoost.

It is important to change the size of the plot because the default one is not readable. The num_trees indicates the tree that should be drawn not the number of trees, so when I set the value to two, I get the second tree generated by XGBoost.

1
2
3
plot_tree(model, num_trees=1)
fig = plt.gcf()
fig.set_size_inches(30, 15)

Did you enjoy reading this article?
Would you like to learn more about software craft in data engineering and MLOps?

Subscribe to the newsletter or add this blog to your RSS reader (does anyone still use them?) to get a notification when I publish a new essay!

Newsletter

Do you enjoy reading my articles?
Subscribe to the newsletter if you don't want to miss the new content, business offers, and free training materials.

Bartosz Mikulski

Bartosz Mikulski

  • Data/MLOps engineer by day
  • DevRel/copywriter by night
  • Python and data engineering trainer
  • Conference speaker
  • Contributed a chapter to the book "97 Things Every Data Engineer Should Know"
  • Twitter: @mikulskibartosz
Newsletter

Do you enjoy reading my articles?
Subscribe to the newsletter if you don't want to miss the new content, business offers, and free training materials.