决策树可视化sklearn

from sklearn import tree
from sklearn.datasets import load_wine
from sklearn.model_selection import train_test_split
import graphviz
import pydotplus

构建数据集

wine = load_wine()
Xtrain,Xtest,Ytrain,Ytest = train_test_split(wine.data,wine.target,test_size=0.3)
print(Xtrain.shape)   #(124, 13)
print(Xtest.shape)    #(54, 13)

构建模型

clf = tree.DecisionTreeClassifier(criterion='entropy')    #实例化
clf = clf.fit(Xtrain,Ytrain)                              #用训练集数据训练模型
score = clf.score(Xtest,Ytest)                            #导入测试集,从接口中调用需要的信息
print(score)                                              #这时结果还在不停的变

画出一棵树

feature_name = ['酒精','苹果酸','灰','灰的碱性','镁','总酚','类黄酮','非黄烷类酚类','花青素','颜色强度','色调','od280/od315稀释葡萄酒','脯氨酸']
dot_data = tree.export_graphviz(clf, out_file='tree.dot',feature_names=feature_name,
                     class_names= ['a','b','c'],   #['琴酒','雪莉','贝尔摩德'],
                     filled=True, rounded=True,
                     special_characters=True)
graph = graphviz.Source(dot_data)
print(graph)

最后在终端Terminal输入:

dot -Tpdf tree.dot -o loan_tree.pdf  生成PDF
dot -Tpng tree.dot -o loan_tree.png  生成png

猜你喜欢

转载自blog.csdn.net/steve_tom/article/details/105867192