可视化sklearn决策树的结构

用graphviz可以直接把树在jupyter notebook里显示出来

import graphviz
from sklearn import tree
model = tree.DecisionTreeRegressor(max_depth=2)
model.fit(X, y)
dot_data = tree.export_graphviz(model, 
                                out_file=None, 
                                feature_names=X.columns,
                                filled=True, 
                                rounded=True,  
                                special_characters=True)  
graph = graphviz.Source(dot_data)
graph

需要通过调用tree.export_graphviz这个方法。

以sklearn自带的数据集iris为例

from sklearn.datasets import load_iris
from sklearn import tree

clf = tree.DecisionTreeClassifier()
iris = load_iris()
clf = clf.fit(iris.data, iris.target)
tree.export_graphviz(clf, out_file='trees.txt')

输出的trees.txt文件里有保存了你的决策树的判断规则。

如果你想把这个规则画出来的话，可以借用这个网站http://webgraphviz.com/，把上面的txt文件里的内容，复制过去就行了，可以得到类似下面的图。