简单Python决策树可视化实例

   本文主要内容参考:http://www.cnblogs.com/pinard/p/6056319.html

   首先需要安装第三方库:第一步是安装graphviz。下载地址在:http://www.graphviz.org/。无论是linux还是windows,装完后都要设置环境变量,将graphviz的bin目录加到PATH,比如我是windows,将C:/Program Files (x86)/Graphviz2.38/bin/加入了PATH。

   第二步是安装python插件graphviz和插件pydotplus。

   本例采用的是sklearn自带的鸢尾花数据集,使用pycharm实现。具体实例如下所示

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Author  : Peidong
# @Site    : 
# @File    : decisiontree.py
# @Software: PyCharm
# 决策树演示
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/'

# 导入数据
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.pdf")


最后所得图像结果为:

简单Python决策树可视化实例