Scikit-learn:决策树可视化

版权声明:本文为博主原创文章,未经博主允许不得转载。 https://blog.csdn.net/pilipala6868/article/details/79963650

平常网上看到的关于决策树的Scikit-learn教程大都是直接给出容器的训练、预测代码:

clf = tree.DecisionTreeClassifier()
clf = clf.fit(X, Y)
clf.predict([...]) 

然而有时候要的不只是预测数据的结果,选择决策树还想要看到训练出来的决策树是什么样的。

于是经过各种搜索,在这里整理一下怎么让生成的决策树可视化。


环境安装、配置

决策树的可视化需要用到graphviz和pydotplus,然而graphviz除了在python里pip安装模块,还需要在电脑中安装graphviz软件。

1、安装graphviz。在官网下载对应版本,比如我是windows,就在Stable 2.38 Windows install packages里下载msi文件后安装到电脑;

2、配置环境变量。根据刚刚安装的路径配置,比如我的是D:\English\Graphviz2.38\bin,要注意配置的路径是bin目录;

3、安装python插件graphviz,这个就直接在命令行里pip install graphviz;

4、安装python插件pydotplus,也是直接pip install pydotplus。


生成可视化文件

首先载入所需的模块:

from sklearn.tree import export_graphviz
from sklearn import tree
import pydotplus 

简单生成一个基本模型:

X = [[0, 0], [1, 1]]
Y = [0, 1]
clf = tree.DecisionTreeClassifier()
clf = clf.fit(X, Y) 

接下来该我们安装的两个模块派上用场:

dot_data = tree.export_graphviz(clf, out_file=None) 
graph = pydotplus.graph_from_dot_data(dot_data)  

最后,可以根据需求生成所需的文件格式:

graph.write_png("tree.png")	# 生成png文件
graph.write_jpg("tree.jpg")	# 生成jpg文件
graph.write_pdf("tree.pdf")	# 生成pdf文件 

如此即可在当前目录中找到我们想要的决策树啦~

阅读更多
想对作者说点什么? 我来说一句

没有更多推荐了,返回首页