调用Python——sklearn包实现决策树

决策树

决策树简介

推荐博客：机器学习——决策树

Python代码——iris数据集

#导入包及数据集
import numpy as np
import pydotplus
from sklearn import tree
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from io import StringIO

#数据集划分
iris=load_iris()
x=np.array(iris.data)
y=np.array(iris.target)
x_train,x_test,y_train,y_test=train_test_split(x,y,test_size=0.2)

#查看划分后的数据维度
x.shape,y.shape
x_train.shape,y_train.shape

#训练决策树模型，参数criterion可选'gini'——CART（默认）或者'entropy'——C4.5
clf=tree.DecisionTreeClassifier(criterion='gini',splitter='best')
clf.fit(x_train,y_train)

#用export_graphviz将树导出为Graphviz格式，用iris.dot文件保存####此步可省略
with open("iris.dot",'w') as f:
    f=tree.export_graphviz(clf,out_file=f)

#########可视化###########
#方法一
dot_data1=tree.export_graphviz(clf,out_file=None,feature_names=iris.feature_names,class_names=iris.target_names,
                              filled=True,rounded=True,special_characters=True)
dot_data1=dot_data1.replace('\n','')####不加这行命令，根节点旁会有一个黑框方格
graph=pydotplus.graph_from_dot_data(dot_data1)
graph.write_pdf("iris.pdf")

#方法二#####不好，有黑点，尝试清楚——失败
dot_data2=StringIO()
tree.export_graphviz(clf,out_file=dot_data2,feature_names=iris.feature_names,class_names=iris.target_names,\
                              filled=True,rounded=True,special_characters=True)
graph=pydotplus.graph_from_dot_data(dot_data2.getvalue())
graph.write_pdf("irisTree.pdf")

可视化

iris.dot

criterion=‘entropy’

iris.pdf

方法二：
irisTree.pdf

可见，根节点旁有个黑框，很丑，尝试去除失败，如果有解决办法，可以在评论区或者私信告诉我~
尝试：

criterion=‘gini’

注：CART决策树一定是二叉树！