机器学习之决策树算法（八）绘制决策树图

以鸢尾花为例绘制决策树图

'''
	使用决策树中的export_graphviz 和绘图的pydotplus来绘制决策树图
'''
# 导入鸢尾花
from sklearn.datasets import load_iris
# 导入数据切分工具
from sklearn.model_selection import train_test_split
# 导入决策树和表格输出函数
from sklearn.tree import DecisionTreeClassifier,export_graphviz
# 导入绘图函数
import pydotplus

# 1.加载数据并做预处理（切分训练集、测试集）。（30分）
data = load_iris()
# 提取数据
x = data.data
y = data.target

# 切分数据
train_x,test_x,train_y,test_y = train_test_split(x,y,test_size=0.3)

# 2.正确建立决策树模型并对数据进行训练，并对预测集进行预测（50分）
model = DecisionTreeClassifier(max_depth=3)
model.fit(train_x,train_y)

# 3.用graphviz和pydotplus画出决策树图。（20分）保存到本地
with open('c.png','wb') as fw:    
    fw.write(pydotplus.graph_from_dot_data(export_graphviz(model)).create_png())

绘图如下