Skip to content

Instantly share code, notes, and snippets.

@hhl60492
Last active April 4, 2019 13:16
Show Gist options
  • Save hhl60492/8bcb9eceaad053c71e9d7d064ad0b013 to your computer and use it in GitHub Desktop.
Save hhl60492/8bcb9eceaad053c71e9d7d064ad0b013 to your computer and use it in GitHub Desktop.
from sklearn.datasets import *
from sklearn import tree
import graphviz
wine = load_wine()
clf = tree.DecisionTreeClassifier() # init the tree
clf = clf.fit(wine.data, wine.target) # train the tree
# export the learned decision tree
dot_data = tree.export_graphviz(clf, out_file=None,
feature_names=wine.feature_names,
class_names=wine.target_names,
filled=True, rounded=True,
special_characters=True)
graph = graphviz.Source(dot_data)
graph.render("wine") # tree saved to wine.pdf
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment