scikit-learn之决策树可视化

Dear 丶 2022-05-23 02:43 192阅读 0赞

平时我们在用机器学习建模时，往往只是用建模去分析数据，得到结论。但有时，我们也需要一些可视化的东西，比如决策树可视化等。  
  在Python的机器学习库scikit-learn中，tree类中的export\_graphviz()函数就能导出树的可视化结果。下面我们将通过一个简单的例子来展示如何将模型建立的决策树可视化。我们使用的数据是位于E盘中log\_reg文件夹下的playTennisTr.csv，数据如下：

![playTennisTr.csv][]

Python代码如下：

# import modules
    import pandas as pd
    from sklearn import tree
    import graphviz
    
    # read data from other places, e.g. csv
    # drop_list: variables that are not used
    def read_data(file_path, drop_list=[]):
        dataSet = pd.read_csv(file_path,sep=',')
        for col in drop_list:
            dataSet = dataSet.drop(col,axis=1)
        return dataSet
    
    # read data in csv format
    file_path = "E://log_reg/playTennisTr.csv"
    dataSet = read_data(file_path)
    target_var = 'PlayTennis'
    
    # decision tree with CART in scikit_learn
    # fit the model by DT in scikit_learn
    clf = tree.DecisionTreeClassifier()
    clf = clf.fit(dataSet.iloc[:, 0:-1], dataSet[target_var])
    
    dot_data = tree.export_graphviz(clf, out_file=None,
                                    feature_names=dataSet.columns[:-1], # 特征名称
                                    class_names=['No', 'Yes'], # 目标变量的类别名
                                    filled=True, rounded=True,
                                    special_characters=True)
    graph = graphviz.Source(dot_data)
    graph.render('example.gv', directory='E:\\log_reg', view=True)
    print('Save example.gv file!\n')

其中，read\_data()函数用来读取指定文件路径的csv文件，目标变量为PlayTennis, 该模型生成的决策树模型为clf. 在tree.export\_graphviz()中，会生成dot文件，在利用graphviz模块的Source()函数可以将其转化为gv文件，如下图：

![生成gv文件][gv]  
  
  对于生成后的gv文件，可以用graphviz软件打开（注意：要将graphviz软件的bin文件添加到系统的环境变量中）。如下：  
  
![graphviz软件打开文件][graphviz]  
  
在graphviz软件中，可以将刚才的gv文件保存为png, jpg, pdf等格式，如下：  
  
![文件保存][70]  
  
我们再去文件夹中查看生成的图片，如下：  
  
![生成的决策树图片][70 1]  
  
  这样我们就能获取建模得到的决策树模型的png, jpg, pdf格式，轻松做到了决策树可视化。  
  本次分享到此结束，欢迎大家交流~~

[playTennisTr.csv]: /images/20220523/955125033dc2468897f7e59808976a18.png
[gv]: /images/20220523/79ee75db18a74a998516936e8ee88242.png
[graphviz]: /images/20220523/c3aca63a428b43438332ebc185067057.png
[70]: /images/20220523/a2f0e20cd9004a6f82a9df599cea271b.png
[70 1]: /images/20220523/863a9750eecb46d78ba53e53b8c5381d.png