Decision Tree: build, prune and visualize it using Python | by Gustavo Hideo | Towards Data Science

Binary Tree is one of the most common and powerful data structures of the computing world. The main reason machine learning engineers like decision trees so much is that it has a low cost to process and it’s really easy to understand (it’s transparent, in opposition to the “black box” from the neural network). It uses a branching method to go over every feature and its different outcomes finding patterns for the target values. For a more detailed explanation of decision trees and different methods of branching, check out this article from Thomas Plapinger.
Let’s get started!
We will be using the Titanic data set from a Kaggle to predict survivors. We will import the data and select some features to work with. ‘Survived’ is our target value.
data = pd.read_csv(‘path_or_link/train.csv’)data = data.loc[:,(‘Survived’,’Pclass’,’Sex’,’Age’,’SibSp’,’Parch’,’Fare’)]
Missing valuesWe have some missing values for ‘age’. To fix it, we could run some predictive analysis and fill up these spots, but that’s not the goal for this article, so we will let it for a future story. For now, let’s just drop those rows with missing values.
Encoding categorical featuresTo keep all the features in numeric format, we will encode the ‘Sex’, that means to transform the values from ‘female’ and ‘male’ to 0 and 1.
data.dropna(inplace=True)#'inplace=True' applies the code to the 'data' object.from sklearn.preprocessing import LabelEncoderle = LabelEncoder()data.Sex = le.fit_transform(data.Sex)
Splitting data set into train and test sets
x = data.iloc[:,1:]   # Second column until the last columny = data.iloc[:,0]    # First column (Survived) is our targetfrom sklearn.model_selection import train_test_split#this function randomly split the data into train and test setsx_train, x_test, y_train, y_test = train_test_split(x, y, test_size=.3)#test_size=.3 means that our test set will be 30% of the train set.
Building the decision tree classifierDecisionTreeClassifier() from sklearn is a good off the shelf machine learning model available to us. It has fit() and predict() methods. The fit() method is the “training” part of the modeling process. It finds the coefficients for the algorithm. Then, the predict() method will use the trained model to make predictions on a new set of data (test set).
dtree = DecisionTreeClassifier()dtree.fit(x_train, y_train)  #train parameters: features and targetpred = dtree.predict(x_test)  #parameter: new data to predict
AccuracyTo determine how good our model is, we use another package from sklearn to give it a measure:
from sklearn.metrics import accuracy_scoreaccuracy_score(y_test, pred)#parameters: targets to be predicted and predictions from new data used before
Accuracy: 0.75It means that we have 75% accuracy in our predictions.
We can plot the tree to see its root, branches, and nodes. We will use a set of new libraries to do so.
from sklearn.tree import export_graphvizfrom sklearn.externals.six import StringIOfrom IPython.display import Image  import pydotplusdot_data = StringIO()export_graphviz(dtree, out_file=dot_data)graph = pydotplus.graph_from_dot_data(dot_data.getvalue())graph.write_png('/tree.png')Image(graph.create_png())
StringIO(): creates an object (empty in this case) to receive a string buffer (the tree will be created first as a string before as an image) in DOT (graph description language) format.export_graphviz(): this function exports the tree in DOT format, generating a representation of the decision tree, which is written into the ‘out_file’.graph_from_dot_data(): will use the DOT object to create the graph.Image(): uses raw data to return an image (png, jpeg, gif).
This tree seems pretty long. Let’s change a couple of parameters to see if there is any effect on the accuracy and also to make the tree shorter.
Criterion: defines what function will be used to measure the quality of a split. The options are “gini” and “entropy”.
Max_depth: defines the maximum depth of the tree. If it’s “none”, the tree will be as long as possible, when all the leaves are pure (risk of overfitting the model).
dtree = DecisionTreeClassifier(criterion='gini')dtree.fit(x_train, y_train)pred = dtree.predict(x_test)print('Criterion=gini', accuracy_score(y_test, pred))dtree = DecisionTreeClassifier(criterion='entropy')dtree.fit(x_train, y_train)pred = dtree.predict(x_test)print('Criterion=entropy', accuracy_score(y_test, pred))
Accuracy results:Criterion=gini: 0.735Criterion=entropy: 0.716
The gini method has a slight improvement over the entropy. Now, let’s check if pruning the tree using max_depth can give us any better results. In the code chunk below, I create a simple function to run our model using different values for max_depth (from 1 to 30) and visualize its results to see how the accuracy differs for each value.
max_depth = []acc_gini = []acc_entropy = []for i in range(1,30): dtree = DecisionTreeClassifier(criterion=’gini’, max_depth=i) dtree.fit(x_train, y_train) pred = dtree.predict(x_test) acc_gini.append(accuracy_score(y_test, pred)) #### dtree = DecisionTreeClassifier(criterion=’entropy’, max_depth=i) dtree.fit(x_train, y_train) pred = dtree.predict(x_test) acc_entropy.append(accuracy_score(y_test, pred)) #### max_depth.append(i)d = pd.DataFrame({‘acc_gini’:pd.Series(acc_gini),  ‘acc_entropy’:pd.Series(acc_entropy), ‘max_depth’:pd.Series(max_depth)})# visualizing changes in parametersplt.plot(‘max_depth’,’acc_gini’, data=d, label=’gini’)plt.plot(‘max_depth’,’acc_entropy’, data=d, label=’entropy’)plt.xlabel(‘max_depth’)plt.ylabel(‘accuracy’)plt.legend()
It seems like in this scenario gini works best for longer trees (as we saw in the previous accuracies), but entropy does a better job for shorter trees and it’s more accurate (near 0.82). So let’s draw the tree with max_depth=7 and criterion=entropy to see how it looks like.
dtree = DecisionTreeClassifier(criterion=’entropy’, max_depth=7)dtree.fit(x_train, y_train)pred = dtree.predict(x_test)accuracy_score(y_test, pred)#Accuracy: 0.814
We went from 0.75 of accuracy to 0.81 just by changing the function to create the branches and limiting the tree to avoid overfitting. Each case is different and it’s a good practice to test different parameters to find the best for your scenario. As you learn new functions and machine learning models, don’t hesitate in reading the documentation. It will give you a broader set of options to customize the model to your needs.
Now, go ahead, use your new set of skills and save the planet by planting a tree.