How to Visualize a Decision Tree Model in Sklearn

2021-02-23

Intro

One large benefit to tree models is that they are easy to intepret and lend themselves to vizualization. You can simple print out a decision tree and follow the log. In this article, we will see how to visualize a tree model using Sklearn.

Vizualizing Tree Models

To vizualize a tree model, we need to do a few steps. We first fit a tree model. We then use the export_graphviz method from the tree module to get dot data. We pass this data to the pydotplus module's graph_from_dot_data function. And finally, we call the write_png function to create our model image.

import pydotplus
from sklearn.tree import DecisionTreeClassifier
from sklearn import datasets
from sklearn import tree

iris = datasets.load_iris()
features = iris.data
target = iris.target

decisiontree = DecisionTreeClassifier()
model = decisiontree.fit(features, target)

dot_data = tree.export_graphviz(decisiontree,
                                out_file = None,
                                feature_names = iris.feature_names,
                                class_names=iris.target_names)


graph = pydotplus.graph_from_dot_data(dot_data)

graph.write_png('tree.png')
GoTea - KoalaTea