Visualizing Decision Tree

scikit-learn.py

# 1. Load the dataset
from sklearn.datasets import load_iris
iris = load_iris()
# 2. Select only the Petal length and Petal width features
#(easier to graph)
X = iris.data[:, 2:]# petal length and width
y = iris.target
# 3. Train our Decision Tree classifier on the Iris Dataset
from sklearn.tree import DecisionTreeClassifier
tree_clf = DecisionTreeClassifier(max_depth=2)
tree_clf.fit(X, y)
# 4. We can visualize the trained decision tree using the
# export_graphviz() method.
from sklearn.tree import export_graphviz
export_graphviz(tree_clf, out_file="tree.dot",
 feature_names=iris.feature_names[2:],
 class_names=iris.target_names,
 rounded=True,
 filled=True)
 
# 5. Convert to png then you can convert this .dot
# file to a variety of formats such as PDF or PNG
# using the dot command- line tool # from the
# graphviz package.
# This command line converts the .dot file to a .png
# image file:
from subprocess import call
call(['dot', '-Tpng', 'tree.dot', '-o', 'tree.png', '-Gdpi=600'])
# 6. Display in python
import matplotlib.pyplot as plt
plt.figure(figsize = (14, 18))
plt.imshow(plt.imread('tree.png'))
plt.axis('off')
plt.show()