Treescope

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

• 展开和折叠已渲染对象的子树，以便您可以关注您关心的模型的部分，
• 自动嵌入任意维数数组和张量的分面可视化，直接嵌入到输出渲染中，以便您可以快速了解它们的形状和值分布，
• 为神经网络模型的部分着色，以强调共享结构，
• 插入"复制路径"按钮，让您可以轻松复制已渲染对象任何部分的路径，
• 自定义可视化策略以支持渲染您自己的数据结构，
• 等等！

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

使用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键！这会打开“往返模式”，并为可视化中的每个类型添加限定名称，使识别对象中的类型更容易。

有关如何使用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产品。