Step 1: Install the libraries
sudo apt-get install graphviz pip install graphviz pip install pydotplus pip install sklearn
pip install pydot pip install pandas
Do the imports
import pydotplus import pandas as pd from sklearn import tree from io import StringIO import pydot
Step 2: Initialize the dataframe
data = [ (0, 5, 0), (1, 6, 0), (2, 7, 1), (3, 8, 1), (4, 9, 1) ] df = pd.DataFrame(data, index=range(5), columns=['x1','x2','y'])
Step 3: Train the decision tree
x_columns = ['x1','x2'] model = tree.DecisionTreeClassifier() trained_model = model.fit(df[x_columns], df['y'])
Step 4: Display the decision tree
Two options
Option A: You want to save the decision tree as a file
dotfile = StringIO() tree.export_graphviz( trained_model, out_file = dotfile, feature_names = x_columns, class_names = ['[y=0]', '[y=1]'], # Ascending numerical order filled = True, rounded = True ) (graph,) = pydot.graph_from_dot_data(dotfile.getvalue()) graph.write_png("tree.png")
This should generate an image named “tree.png” in your current directory
Option B: You want to display the decision tree in your Jupyter notebook
from IPython.display import Image out_file = tree.export_graphviz( trained_model, feature_names = x_columns, class_names = ['[y=0]', '[y=1]'],# Ascending numerical order filled = True, rounded = True ) graph = pydotplus.graph_from_dot_data(out_file) Image(graph.create_png())
In either case this is the tree you should get
References:
https://scikit-learn.org/stable/modules/generated/sklearn.tree.export_graphviz.html