决策树可视化的三种方法

最新推荐文章于 2025-09-29 11:06:06 发布

原创最新推荐文章于 2025-09-29 11:06:06 发布 · 1.6w 阅读

133 ·

CC 4.0 BY-SA版权

文章标签：

#机器学习 #数据可视化 #决策树

本文介绍了三种决策树可视化的技术。第一种是利用sklearn.tree的plot_tree()函数，第二种涉及安装graphviz和使用其Python接口，第三种是通过pydotplus模块进行可视化，该模块同样依赖于Graphviz。每种方法都提供了相应的代码示例。

AI助手已提取文章相关产品：

方法一：直接使用sklearn.tree自带的plot_tree()方法

代码如下：

from sklearn.datasets import load_iris
from sklearn.tree import DecisionTreeClassifier 
from sklearn.tree import plot_tree
from sklearn.model_selection import train_test_split
import matplotlib.pyplot as plt

iris = load_iris()

# 数据拆分
X = iris.data
y = iris.target
X_train,X_test,y_train,y_test = train_test_split(X, y, random_state=0, test_size=1/4)

# 训练模型
dt_model = DecisionTreeClassifier(max_depth=4)
dt_model.fit(X_train, y_train)

# 数据可视化
plt.figure(figsize=(15,9))
plot_tree(dt_model,filled=True,feature_names=iris.feature_names, class_names=iris.target_names)

如图所示：

方法二：使用Graphviz
需要安装
1. graphviz程序，并将安装目录下的bin目录添加到系统path环境变量中
2. graphviz模块, pip install graphviz

代码如下：

from sklearn.datasets import load_iris
from sklearn.tree import DecisionTreeClassifier 
from sklearn.tree import plot_tree
from sklearn.model_selection import train_test_split
from sklearn.tree import export_graphviz
import matplotlib.pyplot as plt
import graphviz


iris = load_iris()

# 数据拆分
X = iris.data
y = iris.target
X_train,X_test,y_train,y_test = train_test_split(X, y, random_state=0, test_size=1/4)

# 训练模型
dt_model = DecisionTreeClassifier(max_depth=4)
dt_model.fit(X_train, y_train)

# 数据可视化
tmp_dot_file = 'decision_tree_tmp.dot'
export_graphviz(dt_model, out_file=tmp_dot_file, feature_names=iris.feature_names, class_names=iris.target_names,filled=True, impurity=False)
with open(tmp_dot_file) as f:
        dot_graph = f.read()
graphviz.Source(dot_graph)

如图所示：

方法三：使用pydotplus模块

pydotplus依赖Graphviz这个绘图库, 它是c开发的, 所以在安装pydotplus之前, 需要首先安装Graphviz

pydotplus是python模块, 它的安装和安装普通的模块一样, 就是使用：

conda install pydotplus

代码如下：

from sklearn.datasets import load_iris
from sklearn.tree import DecisionTreeClassifier 
from sklearn.model_selection import train_test_split
from sklearn.tree import export_graphviz
from IPython.display import Image  
import matplotlib.pyplot as plt
import pydotplus


iris = load_iris()

# 数据拆分
X = iris.data
y = iris.target
X_train,X_test,y_train,y_test = train_test_split(X, y, random_state=0, test_size=1/4)

# 训练模型
dt_model = DecisionTreeClassifier(max_depth=4)
dt_model.fit(X_train, y_train)

# 数据可视化
tmp_dot_file = 'decision_tree_tmp.dot'
export_graphviz(dt_model, out_file=tmp_dot_file, feature_names=iris.feature_names, class_names=iris.target_names,filled=True, impurity=False)
with open(tmp_dot_file) as f:
    dot_graph = f.read()
graph = pydotplus.graph_from_dot_data(dot_graph)
graph.write_pdf('example.pdf')    #保存图像为pdf格式
Image(graph.create_png())   #绘制图像为png格式

结果：