如何获取PyTorch模型的清晰推理计算图？

一、修正Torchviz显示反向图的问题 #

torchviz默认追踪反向传播计算图，要得到正向推理的清晰结构，可按以下操作：

• 用torch.no_grad()包裹推理过程，禁用自动梯度，让计算图仅保留正向节点；
• 调整可视化代码，过滤冗余参数节点：

import torch
from torchviz import make_dot

model = 你的模型实例
input_tensor = torch.randn(你的输入形状)

with torch.no_grad():
    output = model(input_tensor)
# 生成正向计算图，隐藏参数与缓存节点
graph = make_dot(output, params=dict(model.named_parameters()), show_attrs=False, show_saved=False)
graph.render("正向推理图", format="png")

二、解决ONNX算子拆分过细的问题 #

ONNX默认会拆解复杂层为基础算子，要保留原模型层级结构，转换时需做针对性设置：

• 导出ONNX时指定高版本opset、关闭初始化器作为输入，同时启用常量折叠：

torch.onnx.export(
    model,
    input_tensor,
    "model.onnx",
    opset_version=13,
    do_constant_folding=True,
    keep_initializers_as_inputs=False,
    dynamic_axes={
        "input": {0: "batch_size"},
        "output": {0: "batch_size"}
    },
    verbose=False
)

• 用Netron打开导出的ONNX模型，开启界面中的「Group by namespace」或「Collapse subgraphs」选项，即可将拆分的算子重新聚合为原模型的层级结构，得到你需要的清晰图。

三、其他实用替代方案 #

• 自定义模型时，在forward方法中为各层输出命名明确的变量，torchviz可视化时会保留这些命名，提升图的可读性；
• 使用TensorBoard查看层级结构：

from torch.utils.tensorboard import SummaryWriter

writer = SummaryWriter("logs")
writer.add_graph(model, input_tensor)
writer.close()

运行tensorboard --logdir=logs打开面板，即可看到分层的正向推理模型结构。

内容的提问来源于stack exchange，提问作者zjnyly