How to Plot Python XGBdecision tree on Databricks

Question

I am having an issue to plot Xgboost decision tree on databricks. XGboost is installed for python which makes things bit weird when working in this environment.

import xgboost as xgb
from xgboost import plot_importance

gbm=xgb.XGBClassifier().fit(X_train, y_train)
xgb.plot_tree(gbm)

This gives an error:

failed to execute ['dot', '-Tpng'], make sure the Graphviz executables are on your systems' PATH.

I have graphviz installed as a package on databricks side.

Answer 1

I ran into a similiar issue with displaying sklearn decision tree while on databricks.

My solution was to update the source code of the plot_tree function to return the figure in additions to the annotations it normally returns.

Here is an extract of a databricks notebook showing a working example. I imagine that xgboost has a similiar issue/design.

https://github.com/Foley-CJ/Tree_Plot/blob/master/Tree_Plot.ipynb