Treescope

Treescope 是一个交互式 HTML 美化打印器和 N 维数组（“张量”）可视化工具，专为 IPython 笔记本中的机器学习和神经网络研究而设计。它是标准 IPython/Colab 渲染器的直接替代品，并增加了对以下功能的支持：

展开和折叠已渲染对象的子树，以便您专注于模型中您关心的部分，
自动将任意维度数组和张量的分面可视化直接嵌入到输出渲染中，以便您可以快速理解它们的形状及其值的分布，
对神经网络模型的部分进行颜色编码，以突出显示共享结构，
插入“复制路径”按钮，让您轻松复制已渲染对象中任何部分的路径，
自定义可视化策略以支持渲染您自己的数据结构，
等等！

Treescope 最初是作为 Penzai 神经网络库的美化打印器开发的，但它也支持使用其他库开发的神经网络的渲染，包括 Equinox, Flax NNX, 和 PyTorch. 您还可以将其与基本的 JAX 和 Numpy 代码一起使用。

使用 Treescope，您不再需要查看以下内容：

Screenshot of a parameter dictionary without Treescope

您可能正在查看以下内容：

Screenshot of a parameter dictionary using Treescope

这是一个交互式可视化；请尝试点击 ▶ 按钮以展开输出的各个部分！（你也可以在滚动时按住 shift 键，以水平滚动代替垂直滚动。）

Treescope 的文档可在以下地址找到： https://treescope.readthedocs.io.

快速入门

你可以使用以下方式安装 Treescope：

pip install treescope

并使用以下方式导入：

import treescope

要在 IPython 笔记本中使用 Treescope 渲染特定对象，您可以使用 treescope.show，它类似于 print，但会生成丰富的交互式输出。

或者，您可以通过以下方式简单地将 Treescope 配置为笔记本的默认美化打印器：

treescope.register_as_default()

要开启自动数组可视化，您可以运行：

treescope.active_autovisualizer.set_globally(treescope.ArrayAutovisualizer())

或者，如果您希望将 Treescope 设置为默认的格式化打印器，并启用自动数组可视化，您可以简单地运行：

treescope.basic_interactive_setup(autovisualize_arrays=True)

一旦你渲染了一个对象，尝试点击它并按下 r 键！这会开启“往返模式”，并为可视化中的每个类型添加限定名称，从而更容易识别对象中的类型。

[!TIP] 如果 Treescope 的输出过于冗长，或者你正在使用会换行的终端，你可以配置 Treescope 在给定深度处缩写折叠的对象，使用：
treescope.basic_interactive_setup(
    autovisualize_arrays=True,
    abbreviation_threshold=1,  # or a different value
)
你还可以通过覆盖 treescope.abbreviation_threshold 使用 .set_globally 或 .set_scoped 方法手动配置缩写阈值。

有关如何使用 Treescope 的更多信息，请查看 Treescope 文档。

正在寻找一个对 Treescope 的可视化功能提供一流支持的神经网络库？试试 Penzai！

引用

如果你发现 Treescope 对你的研究有用，请考虑引用以下文章（也可在 arXiv 上获取）：

@article{johnson2024penzai,
    author={Daniel D. Johnson},
    title={{Penzai} + {Treescope}: A Toolkit for Interpreting, Visualizing, and Editing Models As Data},
    year={2024},
    journal={ICML 2024 Workshop on Mechanistic Interpretability}
}

这不是一个官方支持的 Google 产品。