机器学习

决策树可视化

2019-06-15  本文已影响0人  凌霄文强
import numpy as np
import pandas as  pd
import pydotplus


from sklearn.datasets import load_iris
iris = load_iris()

from sklearn.cross_validation import train_test_split
# 把数据分为测试数据和验证数据
train_data, test_data, train_target, test_target = train_test_split(iris.data, iris.target, test_size=0.2,
                                                                    random_state=1)
# Model(建模)-引入决策树
from sklearn import tree

# 建立一个分类器
clf = tree.DecisionTreeClassifier(criterion="entropy")
# 训练集进行训练
clf.fit(train_data, train_target)

# 画图方法1-生成dot文件
with open('treeone.dot', 'w') as f:
    dot_data = tree.export_graphviz(clf, out_file=None)
    f.write(dot_data)

# 画图方法2-生成pdf文件
dot_data = tree.export_graphviz(clf, out_file=None, feature_names=clf.feature_importances_,
                                filled=True, rounded=True, special_characters=True)
graph = pydotplus.graph_from_dot_data(dot_data)
###保存图像到pdf文件
graph.write_pdf("treetwo.pdf")
image.png
上一篇下一篇

猜你喜欢

热点阅读