Skip to content

Instantly share code, notes, and snippets.

@JonathanLoscalzo
Forked from hhl60492/tree.py
Created April 4, 2019 13:16
Show Gist options
  • Save JonathanLoscalzo/fa16747e269d360a0a92e1615f75287d to your computer and use it in GitHub Desktop.
Save JonathanLoscalzo/fa16747e269d360a0a92e1615f75287d 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