Hide name of criterion in plot_tree() function

Question

I need the graph from plot_tree() function to be displayed without a criterion value in each node of tree. That is, I want to get a picture like that:

withoun using Paint:) How can I do it?

Answer 1

You can do this by using the impurity=False argument. Here is a reproducible piece of code for you -

from sklearn.datasets import load_iris
from sklearn import tree
import matplotlib.pyplot as plt

#load data
iris = load_iris()

#model training
clf = tree.DecisionTreeClassifier(random_state=0)
clf.fit(iris.data, iris.target)

#plotting
fig = plt.figure(figsize=(10,10))
tree.plot_tree(clf, filled=True, impurity=False)  #<-----------
plt.show()