Graphviz安装及使用：决策树可视化

Python大数据分析与机器学习学习笔记

Graphviz：微信公众号（数学建模与人工智能）
github地址：https://github.com/QInzhengk/Math-Model-and-Machine-Learning

1、决策树模型搭建

# 模型搭建代码汇总 import pandas as pd # 1.读取数据与简单预处理 df = pd.read_excel('员工离职预测模型.xlsx') df = df.replace({ 
   '工资': { 
   '低': 0, '中': 1, '高': 2}}) # 2.提取特征变量和目标变量 X = df.drop(columns='离职') y = df['离职'] # 3.划分训练集和测试集 from sklearn.model_selection import train_test_split X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=123) # 4.模型训练及搭建 from sklearn.tree import DecisionTreeClassifier model = DecisionTreeClassifier(max_depth=3, random_state=123) model.fit(X_train, y_train) 

2、graphviz插件安装与环境变量部署

3、在Python中使用graphviz

from sklearn.tree import export_graphviz import graphviz dot_data = export_graphviz(model, out_file=None, class_names=['0', '1']) graph = graphviz.Source(dot_data) graph.render('决策树可视化') 

那么此时我们可以手动在程序里部署环境变量，只要在代码最上方添加如下两行代码即可。其中添加字母r的原因是用来去除文件路径中反斜杠的特殊含义，另一种方法是用两个反斜杠“\”代替单个斜杠“\”。如果graphviz的安装路径不是下面所示，那么自行修改即可。

import os os.environ['PATH'] = os.pathsep + r'C:\Program Files (x86)\Graphviz2.38\bin' 

完整代码如下：

from sklearn.tree import export_graphviz import graphviz import os # 以下这两行是手动进行环境变量配置，防止在本机的环境变量部署失败 os.environ['PATH'] = os.pathsep + r'C:\Program Files (x86)\Graphviz2.38\bin' dot_data = export_graphviz(model, out_file=None) graph = graphviz.Source(dot_data) graph.render('决策树可视化') 

如果想让可视化的图片里的内容更丰富些，可以在export_graphviz()的括号中增添一些参数，比如增加feature_names参数可以显示特征变量的名称，不过直接使用graphviz识别不了中文，所以我们等会先用英文来做演示；增添class_names参数则可以显示最后的分类结果，因为中文可能会出现乱码问题，所以我们用字符串’0’和’1’来代替’流失’和’非流失’，注意不能写成数字格式的0和1，因为class_names参数中只能设置字符串格式的数据；将filled参数设置为True还可以给可视化的决策树添加颜色，演示代码如下：

# 添加名称（feature_names）和填充颜色（filled=True） dot_data = export_graphviz(model, out_file=None, feature_names=['income', 'satisfication', 'score', 'project_num', 'hours', 'year'], class_names=['0', '1'], filled=True) 

from sklearn.tree import export_graphviz import os # 以下这两行是手动进行环境变量配置，防止在本机的环境变量部署失败 os.environ['PATH'] = os.pathsep + r'C:\Program Files (x86)\Graphviz2.38\bin' 

dot_data = export_graphviz(model, out_file=None, feature_names=X_train.columns, class_names=['不流失', '流失'], rounded=True, filled=True) 

f = open('dot_data.txt', 'w') f.write(dot_data) f.close() 

import re f_old = open('dot_data.txt', 'r') f_new = open('dot_data_new.txt', 'w', encoding='utf-8') for line in f_old: if 'fontname' in line: font_re = 'fontname=(.*?)]' old_font = re.findall(font_re, line)[0] line = line.replace(old_font, 'SimHei') f_new.write(line) f_old.close() f_new.close() 

for line in f_old: if 'fontname' in line: font_re = 'fontname=(.*?)]' old_font = re.findall(font_re, line)[0] line = line.replace(old_font, 'SimHei') f_new.write(line) 

• 黑体 SimHei
• 微软雅黑 Microsoft YaHei
• 新宋体 NSimSun
• 新细明体 PMingLiU
• 细明体 MingLiU
• 仿宋 FangSong
• 楷体 KaiTi

os.system('dot -Tpdf dot_data_new.txt -o 决策树模型.pdf') 

这里是通过调用os.system()系统函数来使用graphviz插件，其中-Tpdf表示生成PDF格式的文件，如果将其改成-Tpng生成PNG格式的图片文件，代码如下：

os.system('dot -Tpng dot_data_new.txt -o 决策树模型.png') 

所有代码整理如下：

from sklearn.tree import export_graphviz import graphviz import os # 以下这两行是手动进行环境变量配置，防止在本机环境的变量部署失败 os.environ['PATH'] = os.pathsep + r'C:\Program Files (x86)\Graphviz2.38\bin' # 生成dot_data dot_data = export_graphviz(model, out_file=None, feature_names=X_train.columns, class_names=['不流失', '流失'], rounded=True, filled=True) # rounded和字体有关，filled设置颜色填充 # 将生成的dot_data内容导入到txt文件中 f = open('dot_data.txt', 'w') f.write(dot_data) f.close() # 修改字体设置，避免中文乱码！ import re f_old = open('dot_data.txt', 'r') f_new = open('dot_data_new.txt', 'w', encoding='utf-8') for line in f_old: if 'fontname' in line: font_re = 'fontname=(.*?)]' old_font = re.findall(font_re, line)[0] line = line.replace(old_font, 'SimHei') f_new.write(line) f_old.close() f_new.close() # 以PNG的图片形式存储生成的可视化文件 os.system('dot -Tpng dot_data_new.txt -o 决策树模型.png') print('决策树模型.png已经保存在代码所在文件夹！') # 以PDF的形式存储生成的可视化文件 os.system('dot -Tpdf dot_data_new.txt -o 决策树模型.pdf') print('决策树模型.pdf已经保存在代码所在文件夹！')