Hide name of criterion in plot_tree() function-CodePudding

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?

I also need to replace the word 'value' to the word 'proba', is it through changing code source only?

CodePudding user response：

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()