sklearn.tree.export_graphviz函数

彬彬侠

已于 2025-01-03 23:13:33 修改

阅读量1k

点赞数 32

分类专栏：智能风控文章标签： sklearn tree graphviz Python 智能风控机器学习

于 2024-12-31 11:43:19 首次发布

本文链接：https://blog.csdn.net/u013172930/article/details/144846579

版权

智能风控专栏收录该内容

34 篇文章

订阅专栏

sklearn.tree.export_graphviz 是 scikit-learn 中用于可视化决策树的重要函数。它会把训练好的决策树模型转化为一个 Graphviz 能够解析的 .dot 文件格式，从而实现决策树的可视化。以下是函数的主要功能与常见用法说明。

1. 基本功能

输入：一个训练好的 DecisionTreeClassifier 或 DecisionTreeRegressor 模型。
输出：一个描述该决策树结构的 .dot 文件（或字符串）。
可视化流程：
1. 使用 export_graphviz 生成 .dot 文件；
2. 使用 Graphviz 工具（如 dot, gv, 或者在线 Graphviz 编辑器）将 .dot 文件转为图片（如 .png 或 .pdf），即可得到树形结构图。

2. 常见用法

下面以 分类决策树 为例，展示如何使用该函数将模型导出为 .dot 格式并最终转换成图像文件。

from sklearn.datasets import load_iris
from sklearn.tree import DecisionTreeClassifier
from sklearn.tree import export_graphviz
import graphviz

# 1. 准备数据 & 训练决策树
iris = load_iris()
X, y = iris.data, iris.target
clf = DecisionTreeClassifier(max_depth=3, random_state=42)
clf.fit(X, y)

# 2. 导出为 .dot 格式并转为 graphviz 可视化对象
dot_data = export_graphviz(
    clf,
    out_file=None,                   # 将结果存储为字符串，而不是写到文件
    feature_names=iris.feature_names,# 特征名称
    class_names=iris.target_names,   # 类别名称
    filled=True,                     # 填充颜色
    rounded=True,                    # 节点框显示圆角
    special_characters=True
)

# 3. 使用 graphviz 渲染为图像
graph = graphviz.Source(dot_data)
graph.render("iris_decision_tree")  # 将输出保存为 PDF 文件 "iris_decision_tree.pdf"

如果你没有安装 graphviz，可以通过 conda install python-graphviz 或者 pip install graphviz 安装，并确保系统上也安装了 Graphviz 可执行程序。

常用参数说明

out_file：
- 默认值为 None。
- 如果指定一个字符串路径（例如 "tree.dot"），则 export_graphviz 会将 .dot 文件直接写入磁盘。
- 如果设置为 None，则函数返回 .dot 格式的字符串，可直接用 Python 对象来进行渲染。
feature_names：
- 用于在决策树可视化的节点上显示具体的特征名称（例如：sepal length (cm)），而不是默认的 X[0], X[1], ...。
class_names：
- 用于展示分类决策树中各类别对应的名称。
- 对于二分类，可以传入 ["class0", "class1"]，对于多分类可以传入列表化的类别名。
filled：
- 如果设置为 True，则各节点根据分类纯度（或者在回归树中根据数值大小）进行不同的颜色填充，便于从颜色深浅快速判断节点的分布。
rounded：
- 如果设置为 True，则使节点和边缘呈圆角矩形显示；如果为 False，则是方形。
special_characters：
- 如果设置为 True，可以在节点标签中使用一些特殊字符（如 ≥, ≤ 等），以使树更直观。
proportion：
- 如果设置为 True，绘制的叶子节点里的样本数会按照数据在该叶节点所占的比例进行缩放；如果为 False，则每个叶子节点的框大小一致。
max_depth：
- 不是 export_graphviz 函数本身的参数，而是决策树的训练参数。若想在可视化时只展示树的前几层，可以使用 max_depth 在训练时对树进行限制，或在可视化时指定 max_depth（scikit-learn 0.21+ 版本之后可以在 export_graphviz 里使用同名参数），以避免树过深导致的可视化不便。

3. 生成 .dot 文件并手动转换图像

若你不想在 Python 内直接渲染，而是想把 .dot 文件生成到磁盘后手动转换为 .png 或 .pdf，可以使用以下方式：

export_graphviz(
    clf,
    out_file="my_tree.dot",
    feature_names=iris.feature_names,
    class_names=iris.target_names,
    filled=True,
    rounded=True,
    special_characters=True
)

# 打开终端或命令行，将 dot 文件转换为 png
# dot -Tpng my_tree.dot -o my_tree.png

这样即可在当前目录下生成一个 my_tree.png 图像文件。

4. 回归树可视化

export_graphviz 同样适用于 DecisionTreeRegressor。只不过在回归情形下，叶子节点呈现的是预测的数值，且颜色深浅往往按照数值大小来区分（如果设置了 filled=True）。

示例：

from sklearn import datasets
from sklearn.tree import DecisionTreeRegressor
from sklearn.tree import export_graphviz

boston = datasets.fetch_openml(name="boston", version=1, as_frame=True, parser="pandas")
X, y= boston.data, boston.target

reg = DecisionTreeRegressor(max_depth=3, random_state=42)
reg.fit(X, y)

export_graphviz(
    reg,
    out_file="reg_tree.dot",
    feature_names=boston.feature_names,
    filled=True,
    rounded=True,
    special_characters=True
)

然后同样用 dot 命令行转换为 .png 或者 .pdf 等格式即可。

5. 常见问题

安装与依赖
- export_graphviz 只是将模型信息导出为 .dot 文件。若需要将 .dot 转成 .png/.pdf，需要在操作系统层面安装 Graphviz（命令行工具）。
- Python 端的包 graphviz 只提供与操作系统中 Graphviz 工具的交互，而不包含其可执行文件。如果缺少可执行文件，可能会报错 Executable "dot" not found...。
树过深导致可视化混乱
- 当决策树过深或特征过多时，导出的图会非常大，导致无法清晰查看。可以：
  - 在训练时设置 max_depth（或其他正则化参数）限制树的复杂度；
  - 或在 export_graphviz 时（在较新版本中）使用 max_depth 参数只可视化前几层。
绘图支持
- 常用的方法是将 out_file=None，返回 .dot 字符串后，用 graphviz.Source 创建可视化对象，或者直接写出 .dot 文件后，用系统命令行转换成图片。
中文或其他非 ASCII 字符支持
- 如果特征名称包含中文（或其他非 ASCII 字符），有时 .dot 文件可能会出现编码问题，可以通过 graphviz.Source(dot_data, encoding='utf-8') 等方式处理，或者手动给 .dot 文件增加 charset="UTF-8" 等声明。
- 也可以尝试通过在特征名中使用 escape sequences 或者为 Graphviz 安装特定字体来解决。

6. 小结

export_graphviz 帮助我们将训练好的决策树模型转化为 Graphviz 可解析的 .dot 文件，实现对决策树的可视化展示。
可视化决策树能够帮助我们更好地理解模型的决策过程和特征的重要程度，是模型可解释性的一个常用手段，尤其在金融风控、信用评分卡建模、医学诊断等对解释性要求较高的场景更是常用。
在实际项目中，若树深度不大（如 max_depth 在 3~5），可视化能够带来非常直观的效果。但对深度过大的树而言，应该合理使用剪枝或 max_depth 等超参数，让树具备更好的泛化性能并易于展示。