sklearn决策树绘制

sklearn决策树绘制
安装Graphviz（https://graphviz.gitlab.io/_pages/Download/Download_windows.html )
一路next，安装Graphviz后，需要设置环境变量：
添加系统环境变量 :建立变量名GRAPHVIZ_DOT ;值为安装的路径，如C:\Program Files (x86)\Graphviz2.34\bin\dot.exe
设置环境变量在用户环境变量添加以下一个变量;建立变量名 GRAPHVIZ_INSTALL_DIR, 值为如C:\Program Files (x86)\Graphviz2.34
在系统环境变量建立变量名PATH中添加Graphviz的bin目录路径，如C:\Program Files (x86)\Graphviz2.34\bin
测试是否安装成功：在cmd中输入dot -version检验是否安装成功

# -*- coding: utf-8 -*-
"""
Created on Fri Nov 23 21:26:09 2018

@author: muli
"""

from sklearn.tree import DecisionTreeClassifier
from sklearn import  datasets
from sklearn import cross_validation

# 可视化需要的包
from sklearn.tree import export_graphviz
from sklearn.externals.six import StringIO
import pydotplus


def load_data():
    '''
    加载用于分类问题的数据集。数据集采用 scikit-learn 自带的 iris 数据集

    :return: 一个元组，用于分类问题。
    元组元素依次为：训练样本集、测试样本集、训练样本集对应的标记、测试样本集对应的标记
    '''
    # scikit-learn 自带的 iris 数据集
    iris=datasets.load_iris() 
    X_train=iris.data
    y_train=iris.target
    # 分层采样拆分成训练集和测试集，测试集大小为原始数据集大小的 1/4
    return cross_validation.train_test_split(X_train, y_train,test_size=0.25,
		random_state=0,stratify=y_train)
    

def test_DecisionTreeClassifier(*data):
    '''
    测试 DecisionTreeClassifier 的用法

    :param data: 可变参数。
    它是一个元组，这里要求其元素依次为：训练样本集、测试样本集、训练样本的标记、测试样本的标记
    :return:  None
    '''
    X_train,X_test,y_train,y_test=data
    # 通过前面的测试，选择最优参数
    clf = DecisionTreeClassifier(criterion='gini',splitter='best',
                                 max_depth=6)
    clf.fit(X_train, y_train)

    print("Training score:%f"%(clf.score(X_train,y_train)))
    print("Testing score:%f"%(clf.score(X_test,y_test)))

    feature_name = ['sepal','sepal_width','petal','petal_width'] 
    target_name = ['setosa','versicolor','virginica']


  
    dot_data = StringIO() 
    export_graphviz(clf,out_file = dot_data,feature_names=feature_name, 
                         class_names=target_name,filled=True,rounded=True, 
                         special_characters=True)
   
    graph = pydotplus.graph_from_dot_data(dot_data.getvalue())
    graph.write_pdf("DecisionTree.pdf")
    print('Visible tree plot saved as pdf.')
    
if __name__=='__main__':
    # 产生用于分类问题的数据集
    X_train,X_test,y_train,y_test=load_data() 
    # 调用 test_DecisionTreeClassifier
    test_DecisionTreeClassifier(X_train,X_test,y_train,y_test)

如图所示：

猜你喜欢