机器学习笔记-基于决策树开发的分类任务（Python）-代码链接在文末

本文链接：https://blog.youkuaiyun.com/m0_72078742/article/details/138205186

本文介绍了如何使用Python的sklearn库加载Iris数据集，训练一个决策树模型，然后通过pydotplus库将其转换为DOT格式并可视化为PNG图像，详细解释了export_graphviz函数的参数及其作用。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

代码解析

导入需要的库

from sklearn.datasets import load_iris
from sklearn import tree
from six import StringIO
import pydotplus

加载sklearn自带的数据集

具体内容可以查看sklearn官网。

iris = load_iris()
clf = tree.DecisionTreeClassifier()
clf = clf.fit(iris.data, iris.target)

决策树格式变换

由于下一步我们希望可视化决策树，这就要使用pydotplus库中的函数，使用前需要先将决策树导出为dot格式。

StringIO对象在python中常用于读写字符串缓冲区，你可以向StringIO对象中写入字符串，然后从中读取字符串，就像操作文件一样。

export_graphviz函数的作用是将决策树的Graphviz表示写入dot_data。

dot_data = StringIO()
tree.export_graphviz(clf, out_file=dot_data,  
                    feature_names=iris.feature_names,  
                    class_names=iris.target_names,  
                    filled=True, rounded=True,  
                    special_characters=True)

export_graphviz函数的各参数如下：

out_file=dot_data 参数指定了输出的文件名或对象。在这个例子中，dot_data 可能是一个文件对象或者一个 StringIO 对象。如果 out_file 是 None，函数会返回一个字符串，包含了 DOT 格式的决策树。

feature_names=iris.feature_names 参数提供了特征的名称。这些名称将用于生成的图形中的节点标签。在这个例子中，使用了 iris 数据集的特征名称。

class_names=iris.target_names 参数提供了目标类的名称。这些名称将用于生成的图形中的叶节点。在这个例子中，使用了 iris 数据集的目标类名称。

filled=True 参数指定了节点是否应该被填充。如果为 True，节点将被填充，颜色表示了分类的主要类别（对于分类问题），值的极值（对于回归问题），或者节点的纯度（对于多输出问题）。

rounded=True 参数指定了节点的形状。如果为 True，节点将被绘制为圆角矩形。

special_characters=True 参数指定了是否应该处理特殊字符。如果为 True，特殊字符将被正确处理，可以在节点的标签中使用。

决策树可视化

使用pydotplus将dot格式转换为png，并存储到程序所在的文件夹中。

graph = pydotplus.graph_from_dot_data(dot_data.getvalue())
graph.write_png("iris.png")

运行结果

代码下载与使用

点击此处免费下载代码使用