scikit-learn之决策树可视化

平时我们在用机器学习建模时，往往只是用建模去分析数据，得到结论。但有时，我们也需要一些可视化的东西，比如决策树可视化等。
在Python的机器学习库scikit-learn中，tree类中的export_graphviz()函数就能导出树的可视化结果。下面我们将通过一个简单的例子来展示如何将模型建立的决策树可视化。我们使用的数据是位于E盘中log_reg文件夹下的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文件，可以用graphviz软件打开（注意：要将graphviz软件的bin文件添加到系统的环境变量中）。如下：

graphviz软件打开文件

在graphviz软件中，可以将刚才的gv文件保存为png, jpg, pdf等格式，如下：

我们再去文件夹中查看生成的图片，如下：

生成的决策树图片

这样我们就能获取建模得到的决策树模型的png, jpg, pdf格式，轻松做到了决策树可视化。
本次分享到此结束，欢迎大家交流~~

scikit-learn之决策树可视化

猜你喜欢