# 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()