dtreeviz：决策树可视化

简介

一个用于决策树可视化和模型解释的 Python 库。决策树是 梯度提升机 和 随机森林(tm) 的基础构建模块，它们可能是处理结构化数据时最流行的两种机器学习模型。在学习这些模型的工作原理以及解释模型时，决策树的可视化可以提供极大的帮助。这些可视化灵感来源于 R2D3 制作的一段教育动画——机器学习的视觉入门。有关我们决策树可视化库及其视觉设计决策的更深入讨论，请参阅 如何可视化决策树。

目前，dtreeviz 支持：scikit-learn、XGBoost、Spark MLlib、LightGBM 和 TensorFlow。安装说明请参见 README.md#Installation。

作者

• Terence Parr，谷歌的技术负责人，2022 年之前曾任旧金山大学的数据科学/计算机科学教授，并于 2012 年担任该大学 数据科学硕士项目 的创始主任。
• Tudor Lapusan
• Prince Grover

此外，Matthew Epland (@mepland) 对代码和可视化进行了大规模的优化和清理。

示例可视化

树形可视化

预测路径解释

叶节点信息

特征空间探索

回归

分类

分类边界

作为一项实用工具，dtreeviz 提供了 dtreeviz.decision_boundaries() 方法，用于展示分类器的一维和二维特征空间，包括表示概率的颜色、决策边界以及被错误分类的样本。值得注意的是，此方法并不局限于树模型，任何实现了 predict_proba() 方法的模型都可以使用。这意味着 scikit-learn 中的任何模型都适用（不过我们也使其能够与定义了 predict() 方法的 Keras 模型配合使用）。由于该方法不专门针对树模型，因此不会使用从 dtreeviz.model() 获取的适配器。更多信息请参阅 classifier-decision-boundaries.ipynb。

有时，查看随超参数变化的动画会很有帮助。如果您查阅笔记本 classifier-boundary-animations.ipynb，您将看到生成如下动画的代码（动画 PNG 文件）：

快速入门

请参阅安装说明，然后查看您所使用的受支持机器学习库的特定笔记本：

• 基于 scikit-learn 的示例（Colab）
• 基于 LightGBM 的示例（Colab）
• 基于 Spark 的示例（Colab）
• 基于 TensorFlow 的示例（Colab）。另请参阅 tensorflow.org 博客中的使用 dtreeviz 可视化 TensorFlow 决策树森林
• 基于 XGBoost 的示例（Colab）
• 适用于任何 scikit-learn 模型的分类器决策边界.ipynb（Colab）
• 更改颜色笔记本（Colab）
• AI 驱动的树分析（scikit-learn）——使用 LLM 进行交互式聊天和解释

为了与这些不同的库进行互操作，dtreeviz 使用一个适配器对象，该对象通过函数 dtreeviz.model() 获取，用于提取可视化所需的相关模型信息。有了这样的适配器对象，您就可以使用相同的编程接口访问 dtreeviz 的所有功能。基本的 dtreeviz 使用流程如下：

1. 导入 dtreeviz 和您的决策树库
2. 获取数据并将其加载到内存中
3. 使用您的决策树库训练分类器或回归器模型
4. 通过
   viz_model = dtreeviz.model(your_trained_model,...) 获取 dtreeviz 适配器模型
5. 调用 dtreeviz 函数，例如
   viz_model.view() 或 viz_model.explain_prediction_path(sample_x)

示例

以下是一个完整的 Python 文件示例，它会在弹出窗口中显示如下树：

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

import dtreeviz

iris = load_iris()
X = iris.data
y = iris.target

clf = DecisionTreeClassifier(max_depth=4)
clf.fit(X, y)

viz_model = dtreeviz.model(clf,
                           X_train=X, y_train=y,
                           feature_names=iris.feature_names,
                           target_name='iris',
                           class_names=iris.target_names)

v = viz_model.view()     # 渲染为 SVG 到内部对象 
v.show()                 # 弹出窗口
v.save("/tmp/iris.svg")  # 可选地保存为 svg

在笔记本中，您可以直接内联渲染，而无需调用 show()。只需调用 view()：

viz_model.view()       # 在笔记本中，直接内联显示

AI 驱动的树分析

启用 AI 集成后，您可以使用 chat() 方法临时询问有关决策树模型的问题。AI 可以访问关于您的树结构、节点和训练数据的全面知识，从而回答以下问题：

• 树结构：整体架构、深度、节点数量、分裂标准以及树的类型（分类/回归）
• 树节点：分裂条件、特征使用情况、节点统计信息、样本分布以及内部节点的纯度指标
• 叶节点：预测结果、置信度分数、样本数量以及类别分布
• 训练数据集：特征统计信息、目标分布以及节点或叶节点内的数据特征

# 创建模型时启用 AI 聊天
viz_model = dtreeviz.model(tree_classifier,
                           X_train=dataset[features], y_train=dataset[target],
                           feature_names=features,
                           target_name=target, class_names=["perish", "survive"],
                           ai_chat=True,
                           ai_model="gpt-4.1-mini",
                           max_history_messages=10)

# 询问关于树的问题
viz_model.chat("请给我一个关于树结构的简短总结？")
viz_model.chat("哪些叶节点的预测置信度最低？")

此外，当启用 ai_chat=True 时，主要的可视化方法（如 view()）会自动在视觉输出旁边包含由 LLM 生成的解释，从而为您提供决策树的图形和自然语言双重解读。

需要运行 pip install dtreeviz[ai] 并将 OpenAI API 密钥设置为 OPENAI_API_KEY 环境变量。更多示例请参阅AI 驱动的树分析笔记本。

安装

如果尚未安装，请在您的系统上安装 anaconda3。

您可能需要确认没有安装由 conda 提供的 graphviz 相关包，因为 dtreeviz 需要 pip 版本；您可以通过以下命令从 conda 环境中移除它们：

conda uninstall python-graphviz
conda uninstall graphviz

要安装（仅限 Python >=3.6），请执行以下操作（在 Windows 上使用 Anaconda Prompt！）：

pip install dtreeviz             # 安装用于 scikit-learn 的 dtreeviz
pip install dtreeviz[xgboost]    # 安装 XGBoost 相关依赖
pip install dtreeviz[pyspark]    # 安装 PySpark 相关依赖
pip install dtreeviz[lightgbm]   # 安装 LightGBM 相关依赖
pip install dtreeviz[tensorflow_decision_forests]   # 安装 TensorFlow 决策树森林相关依赖
pip install dtreeviz[ai]         # 安装 AI 聊天/解释功能（需要 OpenAI API 密钥）
pip install dtreeviz[all]        # 安装所有相关依赖

这也将引入 graphviz Python 库（>=0.9），我们正利用它来处理平台特定的任务。

局限性。 目前只能生成 SVG 文件，这减少了依赖项并大大简化了安装过程。

如果您有任何关于如何让 dtreeviz 在其他平台上更好地工作的建议，请发送电子邮件至 Terence。谢谢！

有关您特定平台的信息，请参阅以下小节。

Mac

请确保已安装最新版本的 Xcode 和命令行工具。如果 Xcode 已经安装，可以在终端中运行 xcode-select --install 来安装这些工具。此外，您还需要签署 Xcode 许可协议，可以通过在终端中运行 sudo xcodebuild -license 来完成。接下来展示的 Homebrew 安装步骤需要编译 Graphviz，因此必须正确配置 Xcode。

您需要 dot 可执行文件来使用 Graphviz。请确保安装的是最新版本（已在 macOS 10.13 和 10.14 上验证）：

brew reinstall graphviz

为确保无误，请从任何 Anaconda 安装中移除 dot，例如：

rm ~/anaconda3/bin/dot

在终端中，以下命令应该能够正常运行：

dot -Tsvg

即它会静默等待输入而不报错。您可以按 Ctrl-C 返回到 shell 提示符。请确认您使用的是通过 Homebrew 安装的正确 dot：

$ which dot
/usr/local/bin/dot
$ ls -l $(which dot)
lrwxr-xr-x  1 parrt  wheel  33 May 26 11:04 /usr/local/bin/dot@ -> ../Cellar/graphviz/2.40.1/bin/dot
$

局限性： Jupyter Notebook 存在一个 bug，无法正确显示 .svg 文件，但 Jupyter Lab 则没有这个问题。

Linux (Ubuntu 18.04)

要获取 dot 可执行文件，请执行以下命令：

sudo apt install graphviz

局限性： view() 方法可以弹出新窗口，并且图像会在 Jupyter Notebook 中内联显示，但在 Jupyter Lab 中则会因解析 SVG XML 出错而无法正常显示。此外，Notebook 中的图像会将我们使用的 Arial 字体替换为其他字体，导致部分文本重叠。在此平台上仅能生成 .svg 文件。

Windows 10

（请务必在所有平台上执行 pip install graphviz，并在 Windows 上通过 Anaconda Prompt 进行安装！）

下载 graphviz-2.38.msi，并更新您的 Path 环境变量。将 C:\Program Files (x86)\Graphviz2.38\bin 添加到用户路径，将 C:\Program Files (x86)\Graphviz2.38\bin\dot.exe 添加到系统路径。由于是 Windows 系统，更新环境变量后可能需要重启。您应在 Anaconda Prompt 中看到如下输出：

(base) C:\Users\Terence Parr>where dot
C:\Program Files (x86)\Graphviz2.38\bin\dot.exe

（请勿使用 conda install -c conda-forge python-graphviz，因为这会导致安装旧版本的 Graphviz Python 库。）

请在 Anaconda Prompt 中验证以下命令是否生效（注意是大写的 -V，而非小写的 -v）：

dot -V

如果无效，则说明您的 Path 设置存在问题。我发现以下测试程序很有用。第一个程序用于检查 Python 是否能找到 dot：

import os
import subprocess
proc = subprocess.Popen(['dot','-V'])
print( os.getenv('Path') )

第二个版本的功能相同，但使用了 Graphviz Python 库的后台支持工具，这也是 dtreeviz 中所采用的方式：

import graphviz.backend as be
cmd = ["dot", "-V"]
stdout, stderr = be.run(cmd, capture_output=True, check=True, quiet=False)
print( stderr )

如果您在运行命令时遇到问题，可以尝试从 https://github.com/xflr6/graphviz/tree/master/graphviz 复制以下文件，并将其放置在 AppData\Local\Continuum\anaconda3\Lib\site-packages\graphviz 文件夹中。同时清空 __pycache__ 目录。

对于 Graphviz Windows 版本 8.0.5 和 Python 接口 v0.18+，您可以执行以下命令：

import graphviz.backend as be
cmd = ["dot", "-V"]
stdout = be.execute.run_check(cmd, capture_output=True，check=True，quiet=False)
print( stdout )

Jupyter Lab 和 Jupyter Notebook 都能很好地显示内联的 .svg 图像。

验证 Graphviz 安装

尝试创建一个名为 t.dot 的文本文件，内容为 digraph T { A -> B }（可以直接复制粘贴到文本编辑器中），然后在终端中运行以下命令：

dot -Tsvg -o t.svg t.dot

这将生成一个简单的 t.svg 文件，应能正常打开。如果 dot 报错，则 dtreeviz 的 Python 代码将无法正常工作。如果找不到 dot，则说明您未正确更新 PATH 环境变量，或 Graphviz 的安装存在其他问题。

局限性

最后，请不要使用 Internet Explorer 查看 .svg 文件，建议使用 Edge 浏览器，效果会更好。我怀疑 IE 是将 .svg 文件渲染为栅格图像而非矢量图像。在此平台上仅能生成 .svg 文件。

在本地安装 dtreeviz

请务必遵循上述安装指南。

为了进行测试，您需要使用 [dev] 附加项安装该库：

pip install dtreeviz[dev]        # 安装开发依赖

在开发过程中，若需将 dtreeviz 库强制更新到本地 egg 缓存，可在 Windows 的 Anaconda Prompt 中执行以下命令：

python setup.py install -f

例如，在 Terence 的机器上，这会添加 /Users/parrt/anaconda3/lib/python3.6/site-packages/dtreeviz-2.3.2-py3.6.egg。

反馈

我们欢迎用户提供关于如何使用 dtreeviz、希望新增哪些功能等方面的信息，可通过 电子邮件（parrt@antlr.org） 或 GitHub 问题页面 与我们联系。

有用资源

• 如何可视化决策树
• 如何解释梯度提升
• 机器学习机制
• R2D3 的动画作品
• 机器学习视觉入门
• fast.ai 面向程序员的机器学习 MOOC
• Stef van den Elzen 的 决策树的交互式构建、分析与可视化
• 类似特征空间可视化的研究，见 迈向用户与计算机在分类任务中的有效协作，SIGKDD 2000
• BigML 决策树之美
• “SunBurst” 树形可视化方法：对用于展示层次结构的空间填充信息可视化技术的评估

许可证

本项目采用 MIT 许可证授权，详情请参阅 LICENSE。

dtreeviz 快速上手指南

dtreeviz 是一个用于决策树可视化和模型解释的 Python 库。它支持 scikit-learn、XGBoost、LightGBM、Spark MLlib 和 TensorFlow Decision Forests 等主流框架，能帮助开发者直观地理解梯度提升树和随机森林等模型的工作原理。

环境准备

• 操作系统：Windows, macOS, Linux
• Python 版本：>= 3.6
• 前置依赖：
  • Graphviz 系统工具：dtreeviz 依赖 Graphviz 进行图形渲染。请务必先在系统中安装 Graphviz 可执行文件，而不仅仅是 Python 包。
    • Ubuntu/Debian: sudo apt-get install graphviz
    • macOS (Homebrew): brew install graphviz
    • Windows: 下载并运行官方安装包 (https://graphviz.org/download/)，并将安装路径添加到系统环境变量 PATH 中。
  • 注意：如果你之前通过 conda 安装了 python-graphviz 或 graphviz，建议先卸载以避免冲突：

    conda uninstall python-graphviz
    conda uninstall graphviz
    

安装步骤

推荐使用 pip 进行安装。国内用户可添加清华源或阿里源加速下载。

1. 安装基础版（仅支持 scikit-learn）：

   pip install dtreeviz -i https://pypi.tuna.tsinghua.edu.cn/simple
   

2. 安装特定框架支持（根据需求选择）：

   # 支持 XGBoost
   pip install dtreeviz[xgboost] -i https://pypi.tuna.tsinghua.edu.cn/simple
   
   # 支持 LightGBM
   pip install dtreeviz[lightgbm] -i https://pypi.tuna.tsinghua.edu.cn/simple
   
   # 支持 PySpark
   pip install dtreeviz[pyspark] -i https://pypi.tuna.tsinghua.edu.cn/simple
   
   # 支持 TensorFlow Decision Forests
   pip install dtreeviz[tensorflow_decision_forests] -i https://pypi.tuna.tsinghua.edu.cn/simple
   

3. 安装所有依赖及 AI 功能（可选）：

   # 安装所有框架支持
   pip install dtreeviz[all] -i https://pypi.tuna.tsinghua.edu.cn/simple
   
   # 安装 AI 聊天与分析功能（需配置 OPENAI_API_KEY 环境变量）
   pip install dtreeviz[ai] -i https://pypi.tuna.tsinghua.edu.cn/simple
   

基本使用

以下是一个基于 scikit-learn 的最简使用示例，展示如何训练模型并可视化决策树。

1. 完整脚本示例

将以下代码保存为 demo.py 运行，将会弹出一个窗口显示决策树图像，并保存为 SVG 文件。

from sklearn.datasets import load_iris
from sklearn.tree import DecisionTreeClassifier
import dtreeviz

# 1. 加载数据
iris = load_iris()
X = iris.data
y = iris.target

# 2. 训练模型
clf = DecisionTreeClassifier(max_depth=4)
clf.fit(X, y)

# 3. 创建 dtreeviz 适配器模型
viz_model = dtreeviz.model(clf,
                           X_train=X, y_train=y,
                           feature_names=iris.feature_names,
                           target_name='iris',
                           class_names=iris.target_names)

# 4. 生成视图
v = viz_model.view()     # 渲染为 SVG 对象
v.show()                 # 弹出窗口显示
v.save("/tmp/iris.svg")  # 可选：保存为文件

2. 在 Jupyter Notebook 中使用

在 Notebook 环境中，可以直接内联显示图像，无需调用 show()。

from sklearn.datasets import load_iris
from sklearn.tree import DecisionTreeClassifier
import dtreeviz

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

viz_model = dtreeviz.model(clf,
                           X_train=iris.data, y_train=iris.target,
                           feature_names=iris.feature_names,
                           target_name='iris',
                           class_names=iris.target_names)

# 直接调用 view() 即可在单元格下方显示图像
viz_model.view()

3. 进阶：AI 驱动的分析（可选）

如果安装了 dtreeviz[ai] 并配置了 OPENAI_API_KEY，可以使用 chat() 方法通过自然语言询问模型结构或获取自动解释。

# 创建模型时启用 ai_chat
viz_model = dtreeviz.model(clf,
                           X_train=X, y_train=y,
                           feature_names=iris.feature_names,
                           target_name='iris',
                           class_names=iris.target_names,
                           ai_chat=True,
                           ai_model="gpt-4.1-mini")

# 询问关于树结构的问题
viz_model.chat("请简要总结这棵树的结构？")
viz_model.chat("哪些叶子节点的预测置信度最低？")

# 可视化时会自动附带 LLM 生成的文字解释
viz_model.view()