决策树的生成(三)

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

决策树的生成(三)_第1张图片

你可能感兴趣的:(机器学习,决策树,算法)