torchvista

一款交互式工具，可在笔记本中直接可视化 PyTorch 模型的前向传播过程——只需一行代码即可实现。支持基于 Web 的笔记本环境，如 Jupyter、Google Colab 和 Kaggle。此外，还允许将可视化结果导出为图像、SVG 和 HTML 格式。

✨ 功能特性

支持拖拽和缩放的交互式图谱

---

可折叠节点，用于分层模块

---

容错性部分可视化（例如形状不匹配时），便于调试

---

点击节点查看参数和属性信息

---

教程与示例

• 分步教程 👉 这里
• 快速 Google Colab 教程 👉 这里（需登录 Colab）
• 查看演示 👉 这里

⚙️ 使用方法

通过 pip 安装：

pip install torchvista

或者通过 conda 安装（更多信息请参阅 torchvista-feedstock）：

conda install -c conda-forge torchvista

在您的 基于 Web 的笔记本（Jupyter、Colab、VSCode Notebook 等）中运行：

import torch
import torch.nn as nn

# 导入 torchvista
from torchvista import trace_model

# 定义您的模块
class LinearModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(10, 5)

    def forward(self, x):
        return self.linear(x)

# 实例化模块和输入张量
model = LinearModel()
inputs = torch.randn(2, 10)

# 开始追踪！
trace_model(model, inputs)

API：trace_model

trace_model(
    model,
    inputs,
    show_non_gradient_nodes=True,
    collapse_modules_after_depth=1,
    forced_module_tracing_depth=None,
    height=800,
    width=None,
    export_format=None,
    show_module_attr_names=False,
    export_path=None,
    show_compressed_view=False,
)

参数	类型	默认值	分类	描述
model	torch.nn.Module	—	追踪	要可视化的模型实例。
inputs	Any	—	追踪	输入到模型中的数据；可以传入单个输入或元组。
show_non_gradient_nodes	bool	True	可视化	显示常量和其他不在梯度图中的节点。
collapse_modules_after_depth	int	1	可视化	初始展开嵌套模块的深度；0 表示全部折叠（节点仍可交互式展开）。
forced_module_tracing_depth	int	None	追踪	要追踪的模块内部的最大深度；None 表示仅追踪用户定义的模块。
height	int	800	可视化	画布高度，单位为像素。
width	int | str	None	可视化	画布宽度；可接受像素或百分比值；省略时默认为可用宽度。
export_format	str	None	导出	可选的导出格式：png、svg 或 html，用于将图谱保存为文件。若未指定，则默认在笔记本中显示图谱。
show_module_attr_names	bool	False	可视化	在可能的情况下，显示模块的属性名称，而非仅显示类名。
export_path	str	None	导出	自定义导出路径。目前仅支持以 HTML 格式进行自定义路径导出。若仅指定文件名，则会在当前工作目录下创建该文件。
show_compressed_view（实验性功能）	bool	False	可视化	通过将具有相同输入输出维度且类型相同的重复节点合并为单个“重复”块来压缩图谱。此功能目前仅识别 Sequential 和 ModuleList 中的重复节点。警告：该功能在大型模型上可能会消耗较多资源。

torchvista 快速上手指南

torchvista 是一款交互式工具，只需一行代码即可在 Notebook（如 Jupyter、Google Colab、Kaggle）中直接可视化 PyTorch 模型的前向传播过程。它支持拖拽缩放、层级折叠、错误容错显示以及节点详情查看，并支持导出为图片、SVG 或 HTML。

环境准备

• 操作系统：Windows / macOS / Linux
• 运行环境：基于 Web 的 Notebook 环境（推荐 Jupyter Lab、Jupyter Notebook、Google Colab 或 VS Code Notebook）
• 前置依赖：
  • Python 3.8+
  • PyTorch (torch)
  • torchvision (可选，视模型需求而定)

安装步骤

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

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

或者使用 conda 安装：

conda install -c conda-forge torchvista

基本使用

在 Notebook 单元格中运行以下代码，即可生成模型结构的交互式可视化图表：

import torch
import torch.nn as nn

# 导入 torchvista
from torchvista import trace_model

# 定义一个简单的模型
class LinearModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(10, 5)

    def forward(self, x):
        return self.linear(x)

# 实例化模型和输入张量
model = LinearModel()
inputs = torch.randn(2, 10)

# 执行追踪并可视化
trace_model(model, inputs)

运行后，你将看到一个可交互的计算图：

• 拖拽与缩放：自由浏览复杂结构。
• 点击节点：查看参数形状、属性等详细信息。
• 折叠/展开：点击层级模块可折叠或展开内部结构。
• 错误调试：即使发生形状不匹配等错误，也能显示部分图谱辅助定位问题。

如需导出图像，可在 trace_model 中指定 export_format 参数（如 "png", "svg", "html"）。