Skip to content

Instantly share code, notes, and snippets.

@Mattamorphic
Created March 29, 2021 16:57
Show Gist options
  • Save Mattamorphic/13072fcf8f1067a8ccdd25637ab8dbf4 to your computer and use it in GitHub Desktop.
Save Mattamorphic/13072fcf8f1067a8ccdd25637ab8dbf4 to your computer and use it in GitHub Desktop.
Visualizing Decision Tree
# 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()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment