如何在sklearn中使用plot_tree绘制决策树时添加标题或图例

plot_tree 本身是基于matplotlib后端渲染绘图的，没有内置标题、图例设置参数，直接调用matplotlib原生的标注接口即可实现需求，不需要修改绘图函数本身。

核心实现逻辑 #

• 首先按照常规流程调用plot_tree绘制决策树，保持matplotlib画布上下文处于激活状态
• 直接调用matplotlib的标题、图例、文本接口添加需要的标注内容即可

1. 添加标题（最适合放置超参数说明） #

标题是放置模型超参数标注最方便的位置，示例代码如下：

import matplotlib.pyplot as plt
from sklearn.tree import plot_tree

# 假设rf为你已经训练完成的随机森林分类器
plt.figure(figsize=(20, 12))
# 示例为绘制随机森林中索引为0的决策树，可根据需求选择要可视化的树
plot_tree(
    rf.estimators_[0],
    feature_names=feature_cols,  # 替换为你的特征名列表
    class_names=class_labels,    # 替换为你的分类标签名列表
    filled=True
)

# 添加带超参数说明的标题
plt.title(
    label="随机森林单棵决策树可视化\n超参数配置："
          f"n_estimators={rf.n_estimators}, max_depth={rf.max_depth}, "
          f"min_samples_split={rf.min_samples_split}, min_samples_leaf={rf.min_samples_leaf}, "
          f"random_state={rf.random_state}",
    fontsize=14,
    pad=20  # 调整标题与绘图区的间距，避免内容遮挡
)

plt.tight_layout()
plt.show()

超参数取值不需要手动硬编码，直接从训练完成的模型实例中读取对应属性即可，避免手动录入出错。

2. 添加自定义图例 #

如果需要补充节点颜色、元素含义的图例，可以自定义图例元素后添加到画布合适位置：

from matplotlib.patches import Patch

# 自定义图例项，颜色需要和plot_tree的filled=True时的节点配色对应
legend_items = [
    Patch(facecolor="#8ab4f8", edgecolor="black", label="多数类为负样本的节点"),
    Patch(facecolor="#f28b82", edgecolor="black", label="多数类为正样本的节点"),
    Patch(facecolor="white", edgecolor="black", label="纯节点/无分类偏向节点")
]

# 放置图例，可通过bbox_to_anchor参数调整位置避免遮挡决策树内容
plt.legend(
    handles=legend_items,
    loc="upper right",
    bbox_to_anchor=(1.18, 1),
    fontsize=12
)

3. 多子图场景适配 #

如果你需要一次性绘制随机森林中的多棵决策树（用子图排布），不要用全局的plt接口，要针对每个子图的ax对象设置标注：

fig, axes = plt.subplots(nrows=2, ncols=2, figsize=(24, 16))
# 绘制4棵决策树示例
for i, ax in enumerate(axes.flat):
    plot_tree(rf.estimators_[i], feature_names=feature_cols, filled=True, ax=ax)
    # 单独给每个子图加标题
    ax.set_title(f"随机森林第{i+1}棵决策树\n树实际最大深度：{rf.estimators_[i].get_depth()}", fontsize=12)

plt.tight_layout()
plt.show()

如果需要在图内空白位置添加自由文本备注，直接调用plt.text(x, y, "备注内容")传入对应坐标位置即可，位置调整灵活度很高。

内容的提问来源于stack exchange，提问作者NoTisan