torchview

Torchview 提供以可视化图的形式展示 PyTorch 模型的功能。可视化内容包括张量、模块、torch.functions 以及输入输出形状等信息。

Keras 中 plot_model 的 PyTorch 版本（且功能更丰富）

支持 PyTorch 1.7 及以上版本。

有用特性

安装

首先，你需要安装 graphviz，

# uv
uv add graphviz

# pip
pip install graphviz

为了让 graphviz 的 Python 接口正常工作，你的系统中需要有可用的 dot 布局命令。如果尚未安装，建议根据你的操作系统运行以下命令：

基于 Debian 的 Linux 发行版（如 Ubuntu）：

apt-get install graphviz

Windows：

choco install graphviz

macOS：

brew install graphviz

更多详情请参阅 这里

然后，继续使用 uv（或 pip）安装 torchview：

# uv
uv add torchview

# pip
pip install torchview

或者如果你希望通过 conda 安装：

conda install -c conda-forge torchview

如果你想获得最新版本，可以直接从仓库安装：

# uv
uv add git+https://github.com/mert-kurttutan/torchview.git

# pip
pip install git+https://github.com/mert-kurttutan/torchview.git

使用方法

from torchview import draw_graph

model = MLP()
batch_size = 2
# device='meta' -> 不会为可视化消耗内存
model_graph = draw_graph(model, input_size=(batch_size, 128), device='meta')
model_graph.visual_graph

笔记本示例

更多示例，请参阅下方的 Colab 笔记本：

入门笔记本：

计算机视觉模型：

NLP 模型：

注意： 输出的 graphviz 可视化图像会以所需尺寸返回。但在某些情况下，由于图像较大且 VSCode 默认启用 SVG 渲染，部分形状可能会被裁剪。为了解决这个问题，建议你运行以下代码：

import graphviz
graphviz.set_jupyter_format('png')

此问题在其他平台（如 JupyterLab 或 Google Colab）上不会出现。

支持的功能

• 几乎所有模型，包括 RNN、Sequential、跳跃连接、Hugging Face 模型
• 支持 meta 张量 -> 不消耗内存（适用于超大型模型）（PyTorch 1.13 及以上版本）
• 显示张量之间的操作（除了模块调用之外）
• 展开/折叠功能。递归使用的模块可以进行可视化展开，如下所示。
• 多种输入/输出类型，例如嵌套数据结构（字典、列表等）、Hugging Face 分词器输出等

文档

def draw_graph(
    model: nn.Module,
    input_data: INPUT_DATA_TYPE | None = None,
    input_size: INPUT_SIZE_TYPE | None = None,
    graph_name: str = 'model',
    depth: int | float = 3,
    device: torch.device | str | None = None,
    dtypes: list[torch.dtype] | None = None,
    mode: str | None = None,
    strict: bool = True,
    expand_nested: bool = False,
    graph_dir: str | None = None,
    hide_module_functions: bool = True,
    hide_inner_tensors: bool = True,
    roll: bool = False,
    show_shapes: bool = True,
    save_graph: bool = False,
    filename: str | None = None,
    directory: str = '.',
    **kwargs: Any,
) -> ComputationGraph:
    '''返回输入 PyTorch 模块的可视化表示，以 ComputationGraph 对象的形式。ComputationGraph 对象包含：

    1) 根节点（通常是输入张量的张量节点），这些节点连接到在前向传播过程中记录的 PyTorch 模块计算图中的所有其他节点。

    2) graphviz.Digraph 对象，它包含了 PyTorch 模块计算图的可视化表示。该图展示了模块/模块层次结构、PyTorch 函数、形状以及在前向传播过程中记录的张量等信息，具体示例请参阅文档和 Colab 笔记本。


    参数：
        model (nn.Module):
            要进行可视化表示的 PyTorch 模型。

        input_data (包含 torch.Tensor 的数据结构):
            模型 forward 方法的输入。如果是多个参数，请将其包装在列表中；如果是关键字参数，则包装在字典中。

        input_size (尺寸序列):
            输入数据的形状，以 List/Tuple/torch.Size 的形式提供。
            （数据类型必须与模型输入匹配，默认为 FloatTensor）。
            默认值：None

        graph_name (str):
            graphviz.Digraph 对象的名称。同时也是图形可视化文件的默认名称。
            默认值：'model'

        depth (int):
            可视化中显示的节点深度上限。深度是指模块或张量在模块层次结构中的嵌套程度。例如，主模块的深度为 0，而其子模块的深度为 1，依此类推。
            默认值：3

        device (str 或 torch.device):
            用于放置输入张量的设备。如果 PyTorch 检测到 CUDA，则默认使用 GPU，否则使用 CPU。
            默认值：None

        dtypes (torch.dtype 列表):
            如果提供了输入尺寸，则使用这些数据类型来设置输入张量的类型。

        mode (str):
            用于前向传播的模型模式。如果没有指定，则默认使用评估模式。
            默认值：None

        strict (bool):
            如果为真，graphviz 可视图不允许节点之间有多条边。例如，当存在从模块节点到模块节点的张量时，就会出现多条边，此时会隐藏这些张量。
            默认值：True

        expand_nested(bool):
            如果为真，则以虚线边框显示嵌套模块。

        graph_dir (str):
            设置可视化图的方向。
            'TB' -> 从上到下
            'LR' -> 从左到右
            'BT' -> 从下到上
            'RL' -> 从右到左
            默认值：None -> TB

        hide_module_function (bool):
            决定是否隐藏模块的 PyTorch 函数。有些模块仅由 PyTorch 函数组成（没有子模块），例如 nn.Conv2d。
            True => 不将模块函数包含在 graphviz 中
            False => 将模块函数包含在 graphviz 中
            默认值：True

        hide_inner_tensors (bool):
            内部张量是指计算图中除输入和输出张量外的所有张量。
            True => 不在 graphviz 中显示内部张量
            False => 在 graphviz 中显示内部张量
            默认值：True

        roll (bool):
            如果为真，则展开递归模块。
            默认值：False

        show_shapes (bool):
            True => 显示张量、输入和输出的形状
            False => 不显示
            默认值：True

        save_graph (bool):
            True => 保存 graphviz 图的输出文件
            False => 不保存
            默认值：False

        filename (str):
            用于存储 dot 语法表示和 graphviz 图像文件的文件名。默认为 graph_name。

        directory (str):
            存放 graphviz 输出文件的目录。
            默认值：.

    返回：
        包含输入 PyTorch 模型可视化表示的 ComputationGraph 对象，形式为 graphviz Digraph 对象。
    '''

示例

递归网络的折叠版本

from torchview import draw_graph

model_graph = draw_graph(
    SimpleRNN(), input_size=(2, 3),
    graph_name='RecursiveNet',
    roll=True
)
model_graph.visual_graph

显示/隐藏中间（隐藏）张量和函数

# 显示内部张量和函数
model_graph = draw_graph(
    MLP(), input_size=(2, 128),
    graph_name='MLP',
    hide_inner_tensors=False,
    hide_module_functions=False,
)

model_graph.visual_graph

ResNet / 跳跃连接 / 对 Torch 操作的支持 / 嵌套模块

import torchvision

model_graph = draw_graph(resnet18(), input_size=(1,3,32,32), expand_nested=True)
model_graph.visual_graph

待办事项

• 显示模块参数信息
• 支持图神经网络 (GNN)
• 支持 GNN 的无向边
• 支持基于 Torch 的函数[^1]

[^1]: 这里，基于 Torch 的函数指的是仅使用 Torch 函数和模块的任何函数。这比模块的概念更为广泛。

贡献

我们非常欢迎所有的问题和拉取请求！

贡献者须知：

• torchview 正在使用最新版本的 Python 积极开发。
• 更改应与 Python 3.9 向后兼容，并遵循 Python 对旧版本的生命周期支持政策。

设置（推荐使用 uv）

uv sync --group dev --group torch
uv run prek install

设置（pip / venv）

python -m venv .venv
source .venv/bin/activate
pip install -e .
# 根据需要安装开发/测试工具（例如：pytest、prek 等）
pip install pytest prek torch torchvision torchtext
prek install

运行测试

# uv
uv run pytest
uv run pytest --overwrite
uv run pytest --no-output

# pip / venv
pytest
pytest --overwrite
pytest --no-output

参考文献

• 与输入处理和验证相关的部分均源自或受 torchinfo 仓库的启发！！。
• 许多软件相关部分（例如测试）也源自或受 torchinfo 仓库的启发。因此，非常感谢 @TylerYep！！！
• 构建可视化图的算法得益于 __torch_function__ 以及对 torch.Tensor 的子类化。在此向所有开发这一 API 的人表示衷心的感谢！！。

torchview 快速上手指南

torchview 是一个用于可视化 PyTorch 模型的工具，能够生成包含张量、模块、函数操作以及输入/输出形状信息的计算图。它类似于 Keras 的 plot_model，但专为 PyTorch 设计。

环境准备

系统要求

• Python: 3.9+
• PyTorch: $\geq$ 1.7 (若需使用 meta 设备节省内存，建议 PyTorch $\geq$ 1.13)

前置依赖

在使用 Python 包之前，必须在系统中安装 Graphviz 软件及其命令行工具（dot），否则无法渲染图像。

请根据您的操作系统运行以下命令安装 Graphviz：

• Debian/Ubuntu:

  apt-get install graphviz
  

• Windows (需先安装 Chocolatey):

  choco install graphviz
  

• macOS:

  brew install graphviz
  

注意：仅安装 Python 包 graphviz 是不够的，必须确保系统环境变量中可以运行 dot 命令。

安装步骤

您可以选择以下任意一种方式安装 torchview。推荐使用国内镜像源以加速下载。

方式一：使用 pip (推荐)

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

同时确保已安装 python 接口包（通常上述命令会自动处理，若报错可单独执行）：

pip install graphviz

方式二：使用 Conda

conda install -c conda-forge torchview

方式三：安装最新开发版

pip install git+https://github.com/mert-kurttutan/torchview.git

基本使用

以下是最简单的使用示例，展示如何绘制一个 MLP 模型的结构图。

代码示例

from torchview import draw_graph
import torch.nn as nn

# 定义一个简单的模型示例
class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Linear(64, 10)
        )
    
    def forward(self, x):
        return self.net(x)

model = MLP()
batch_size = 2

# 绘制模型图
# device='meta' 表示使用元张量，可视化过程不消耗显存/内存
model_graph = draw_graph(model, input_size=(batch_size, 128), device='meta')

# 显示图形
model_graph.visual_graph

参数说明

• model: 需要可视化的 PyTorch nn.Module 实例。
• input_size: 输入数据的形状（不含批次维度的元组或包含批次维度的元组均可，如上例）。
• device='meta': 关键技巧。设置为 'meta' 可以在不实际分配内存的情况下构建计算图，非常适合超大模型的可视化。

在 VSCode 中的注意事项

如果在 VSCode 中查看生成的 SVG 图片出现裁剪问题，建议在 Notebook 或脚本开头添加以下代码强制使用 PNG 格式渲染：

import graphviz
graphviz.set_jupyter_format('png')