how to add a title (or legend) when plotting a decision tree in sklearn (with plot_tree)

Viewed 107

I have the following random forest classifier graph that I get with plot_tree of sklearn :

enter image description here

But I would like to add a legend (or title) somewhere with the hyperparameter values that I used to train my model. The problem is that such an argument is not present when using plot_tree, how would you handle it to add these informations to the above plot?

1 Answers

plot_tree plots on the current matplotlib.pyplot axes by default. From there you can make use of matplotlib functionality. If you want, you can use the ax parameter to plot onto a specified axes object instead; in the below example you don't really need to call the figure and axes lines, but it might be helpful depending on how you end up decorating the plot.

from matplotlib import pyplot as plt
plt.figure()
ax = plt.axes()
tree.plot_tree(clf, ax=ax)
plt.suptitle("Decision Tree")
hp_text = '\n'.join(f"{k}: {v}" for k, v in clf.get_params().items())
plt.text(1.0, 0.25, hp_text, transform=plt.gca().transAxes);

iris dataset decision tree, with parameters as subtitle

Related