保存决策树模型

"""
步骤:
graphviz下载地址:https://graphviz.gitlab.io/_pages/Download/Download_windows.html
1)安装graphviz,并设置环境变量:将graphviz的bin目录加到PATH
2)安装pydotplus ,打开anaconda prompt, 输入命令 :conda install -c conda-forge pydotplus
3)在anaconda prompt中, 输入命令 :pip install graphviz ,安装graphviz插件
"""

from sklearn.datasets import load_iris
from sklearn import tree
import pydotplus
import os

# 导入路径
# os.environ["PATH"] += os.pathsep + 'C:/Program Files (x86)/Graphviz2.38/bin/'
# os.environ["PATH"] += os.pathsep + 'C:/Users/Dominic Jiang/graphviz-2.38/release/bin'
# 导入数据
iris = load_iris()
# 构建模型
clf = tree.DecisionTreeClassifier()
clf = clf.fit(iris.data, iris.target)

# 保存模型
with open("iris.dot", 'w') as f:
    f = tree.export_graphviz(clf, out_file=f)

# 画图,保存到pdf文件
# 设置图像参数
dot_data = tree.export_graphviz(clf, out_file=None,
                                feature_names=iris.feature_names,
                                class_names=iris.target_names,
                                filled=True, rounded=True,
                                special_characters=True)
graph = pydotplus.graph_from_dot_data(dot_data)
# 保存图像到pdf文件
graph.write_pdf("iris_color.pdf")  # 1.利用graph写入文件

a = open("iris.png", mode="wb")  # 2.利用I/o写文件
a.write(graph.create_png())

os.system("dot -Tpdf iris.dot -o iris_w.pdf")  # 3.利用dot系统命令写文件

你可能感兴趣的:(机器学习)