from sklearn.datasets import load_iris
iris = load_iris()
data = iris.data
target = iris.target
from sklearn.model_selection import train_test_split
x_train, x_test, y_train, y_test = train_test_split(data, target, test_size = 0.3)
from sklearn.tree import DecisionTreeClassifier, export_graphviz
def train():
model = DecisionTreeClassifier(criterion = 'gini')
model.fit(x_train, y_train)
print("在测试集上的准确率为:",model.score(x_test, y_test))
return model
model = train()
import graphviz
feature_names = iris.feature_names
dot_data = export_graphviz(model, out_file=None,
feature_names=feature_names,
filled=True,rounded=True,
special_characters=True)
graph = graphviz.Source(dot_data)
graph.render('iris')