Code snippet for decision tree and plot the decision tree to file

decision_tree.py

tree_clf = tree.DecisionTreeClassifier(random_state=SEED, max_depth=3)
_ = myUtilityFunction.prediction_evaluation(tree_clf, X_train, X_test, 
                                            y_train, y_test, X_train.columns, 
                                            'features')

dot_data = tree.export_graphviz(tree_clf, out_file=None, 
                                feature_names=X_train.columns,  
                                class_names=['No', 'Yes'],  
                                filled=True, rounded=True,  
                                special_characters=True)  
graph = graphviz.Source(dot_data)  
graph.format = 'png'
graph.render('../img/output/model_decision tree_employee attrition', view=False)