平常网上看到的关于决策树的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文件 

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