Linux下python决策树可视化

2019年5月25日 21:16 by wst

算法实践

---

在决策树算法调试过程中，有时需要查看每个节点的情况，此时需要看到树的样子，以找到问题的根源。

 

安装必要的工具

• ubuntu系统

sudo apt-get install graphviz

• centos系统

sudo yum install graphviz

演示代码

"""
Author: deepinwst
Email: wanshitao@donews.com
Date: 18-12-6 下午2:09
"""

from sklearn.datasets import load_iris
from sklearn import tree
import pydot


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)

# 注意，方法一需要先执行这个命令
# dot -Tpdf iris.dot -o iris.pdf

# 方法二：
dot_data = tree.export_graphviz(clf, out_file=None)
graph = pydot.graph_from_dot_data(dot_data)
# graph[0].write_pdf("iris.pdf")

# 方法三：
from IPython.display import Image
import pydotplus

# 需要在notebook中运行
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)
Image(graph.create_png())

完整实例

系统环境(Python 3.6.8 - 64bit)

Distributor ID: Deepin
Description:    Deepin 15.10
Release:    15.10
Codename:   unstable

deepin是类似ubuntu的系统。

注意：此实例是在notebook中运行的。

在使用过程中，若缺少什么包，直接pip install packagename 即可。packagename为你要装的包名。

from sklearn.externals import joblib
from IPython.display import Image
import pydotplus
from sklearn import tree

# 载入训练好的决策树模型
model = joblib.load("age_model.pkl")
dot_data = tree.export_graphviz(model, out_file=None,
                                filled=True, rounded=True,
                                special_characters=True)
graph = pydotplus.graph_from_dot_data(dot_data)
Image(graph.create_png())

 

后话：至于在Windows下怎么安装可视化工具，以及可视化方法，参见：https://blog.csdn.net/pilipala6868/article/details/79963650

---