sachinhosmani/torchvista

What it solves

它提供了一种在基于网页的 notebook（如 Jupyter、Colab 与 Kaggle）内直接可视化 PyTorch 模型前向传播的方法，只需一行代码，即可更轻松地理解与调试模型架构。

How it works

通过使用 trace_model 函数，工具会在给定特定输入时追踪模型的执行，并生成交互式图形。此图形允许用户交互式探索模型结构，包括嵌套模块与参数信息。

Who it’s for

PyTorch 开发者与研究人员，需要在开发与调试过程中可视化模型的数据流与架构。

Highlights

• Interactive Visualization：支持拖拽缩放与可折叠节点，以呈现层级模块。
• Error-Tolerant：即使出现形状不匹配等错误，也能提供部分可视化，帮助调试。
• Detailed Node Info：点击节点即可显示参数和属性信息。
• Export Options：图形可导出为 PNG、SVG 或 HTML 文件。
• Compressed View：实验性功能，将 Sequential 或 ModuleList 中相同类型的重复节点合并为单一块。