Visualize Decision Tree with Python Sklearn Library

Decision tree visualization using GraphViz

In this post, you will learn about different techniques you can use to visualize decision tree (a machine learning algorithm) using Python Sklearn (Scikit-Learn) library. The python code example would use Sklearn IRIS dataset (classification) for illustration purpose. The decision tree visualization would help you to understand the model in a better manner.

The following are two different techniques which can be used for creating decision tree visualisation:

  • Sklearn tree class (plot_tree method)
  • Graphviz library

Sklearn Tree Class for Visualization

In this section, you will see the code sample for creating decision tree visualization using Sklearn Tree method plot_tree method. Sklearn IRIS dataset is used for training the model.

import matplotlib.pyplot as plt
from sklearn import datasets
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier
from sklearn import tree
# Load the IRIS dataset
iris = datasets.load_iris()
X =[:, 2:]
y =
# Create train and test split
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=1, stratify=y)
# Train the model using DecisionTree classifier
clf_tree = DecisionTreeClassifier(criterion='gini', max_depth=4, random_state=1), y_train)
# Plot the decision tree
fig, ax = plt.subplots(figsize=(10, 10))
tree.plot_tree(clf_tree, fontsize=10)

Here is how the decision tree would look like:

Decision tree visualization using Sklearn.tree plot_tree method
Fig 1. Decision tree visualization using Sklearn.tree plot_tree method

GraphViz for Decision Tree Visualization

In this section, you will learn about how to create a nicer visualization using GraphViz library. Here are the set of libraries such as GraphViz, PyDotPlus which you may need to install (in order) prior to creating the visualization. PyDotPlus converts dot data files into a decision tree image file.

pip install graphviz
pip install pyparsing
pip install pydotplus

Here is the code which can be used for creating visualization. It uses the instance of decision tree classifier, clf_tree, which is fit in the above code. Note some of the following in the code:

  • export_graphviz function of Sklearn.tree is used to create the dot file
  • Function, graph_from_dot_data is used to convert the dot file into image file.
from pydotplus import graph_from_dot_data
from sklearn.tree import export_graphviz

dot_data = export_graphviz(clf_tree, filled=True, rounded=True,
                                    feature_names=['petal length',
                                                'petal width'],
graph = graph_from_dot_data(dot_data)

Here is how the tree visualization looks like. Note the difference between the tree visualization created using GraphViz (fig 2) and without using GraphViz (fig 1). Note some of the following in the tree drawn below:

  • Root node splits the training dataset (105) into two child nodes with 35 and 70 data points.
  • Left child node can be said as a pure or homogenous node as it has all the data points belonging to Setosa class. Thus, Gini impurity score is 0.
  • Right child node is split further into two child nodes. The left child node results in the pure data set belonging to Versicolor class with Gini impurity as 0.
Decision tree visualization using Graphviz with max_depth as 4
Fig 2. Decision tree visualization using Graphviz (Max depth = 4)

Change the max_depth of the tree as 3 and this is how the tree will look like.

Decision tree visualization using Graphviz with max_depth as 3
Fig 3. Decision tree visualization using Graphviz (Max depth = 3)

In this post, you learned about how to create a visualization diagram of decision tree using two different techniques (Sklearn.tree plot_tree method) and GraphViz method.

Ajitesh Kumar
Follow me

Ajitesh Kumar

I have been recently working in the area of Data Science and Machine Learning / Deep Learning. In addition, I am also passionate about various different technologies including programming languages such as Java/JEE, Javascript, Python, R, Julia etc and technologies such as Blockchain, mobile computing, cloud-native technologies, application security, cloud computing platforms, big data etc. I would love to connect with you on Linkedin.
Posted in Data Science, Machine Learning. Tagged with , .

Leave a Reply

Your email address will not be published. Required fields are marked *

Time limit is exhausted. Please reload the CAPTCHA.