Visualize a Decision Tree with Sklearn

Step 1: Install the libraries
sudo apt-get install graphviz

pip install graphviz
pip install pydotplus
pip install sklearn
pip install pydot pip install pandas

Do the imports

import pydotplus
import pandas as pd
from sklearn import tree
from io import StringIO
import pydot
Step 2: Initialize the dataframe
data = [ 
    (0, 5, 0), 
    (1, 6, 0), 
    (2, 7, 1), 
    (3, 8, 1), 
    (4, 9, 1)
]
df = pd.DataFrame(data, index=range(5), columns=['x1','x2','y'])
Step 3: Train the decision tree
x_columns = ['x1','x2']

model = tree.DecisionTreeClassifier()
trained_model = model.fit(df[x_columns], df['y'])
Step 4: Display the decision tree

Two options

Option A: You want to save the decision tree as a file

dotfile = StringIO()

tree.export_graphviz(
    trained_model,  
    out_file        = dotfile,
    feature_names   = x_columns, 
    class_names     = ['[y=0]', '[y=1]'], # Ascending numerical order
    filled          = True,
    rounded         = True
)

(graph,) = pydot.graph_from_dot_data(dotfile.getvalue())
graph.write_png("tree.png")

 

This should generate an image named “tree.png” in your current directory

Option B: You want to display the decision tree in your Jupyter notebook

from IPython.display import Image

out_file = tree.export_graphviz(
    trained_model,
    feature_names   = x_columns,
    class_names     = ['[y=0]', '[y=1]'],# Ascending numerical order
    filled          = True,
    rounded         = True
)
graph = pydotplus.graph_from_dot_data(out_file)
Image(graph.create_png())

In either case this is the tree you should get

 

References:

https://scikit-learn.org/stable/modules/generated/sklearn.tree.export_graphviz.html