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