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文件，如下图：