仅通过前向传播提取PyTorch模型计算图的可行方案？

假设有如下CoolCNN模型：

import torch
import torch.nn as nn

class CoolCNN(nn.Module):
    def __init__(self):
        super(CoolCNN, self).__init__()
        self.initial_conv = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, padding=1)
        self.parallel_conv = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, padding=1)
        self.secondary_conv = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, padding=1)
        self.max_pool = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
        self.fully_connected1 = nn.Linear(32 * 8 * 8, 128)
        self.output_layer = nn.Linear(128, 10)

    def forward(self, x):
        main_path = self.max_pool(torch.relu(self.initial_conv(x)))
        parallel_path = self.max_pool(torch.relu(self.parallel_conv(x)))
        x = (main_path + parallel_path) / 2
        x = self.max_pool(torch.relu(self.secondary_conv(x)))
        x = x.view(-1, 32 * 8 * 8)
        x = torch.relu(self.fully_connected1(x))
        x = self.output_layer(x)
        return x

该模型的树状计算结构如下：

CoolCNN
└── Forward Pass
    ├── parallel_path
    │   ├── parallel_conv (Conv2d)
    │   ├── ReLU Activation
    │   └── max_pool (MaxPool2d)
    │
    ├── main_path
    │   ├── initial_conv (Conv2d)
    │   ├── ReLU Activation
    │   └── max_pool (MaxPool2d)
    │
    ├── Average main_path and parallel_path
    │
    ├── secondary_conv (Conv2d)
    ├── ReLU Activation
    └── max_pool (MaxPool2d)
    │
    ├── Flatten the Tensor
    │
    ├── fully_connected1 (Linear)
    ├── ReLU Activation
    │
    └── output_layer (Linear)

希望仅通过模型的前向传播过程提取这种树状计算结构，完全不依赖反向传播。torchviz等库依赖反向传播生成计算图，不符合需求；forward hooks只能获取节点调用顺序（拓扑排序），无法还原唯一的树状结构。请问是否存在可行的方法？

有两种可靠的方法可以实现需求，均完全基于前向传播，无需反向传播：

方案1：使用PyTorch官方torch.fx模块 #

torch.fx是PyTorch内置的静态图捕获工具，能直接通过前向传播生成模型的中间表示（IR），包含所有操作节点的依赖关系，可准确还原分支结构。

实现步骤 #

1. 用torch.fx.symbolic_trace追踪模型，生成包含计算图的GraphModule
2. 遍历图中节点，每个节点的args属性记录了它的输入依赖，target属性记录了操作类型或模块名
3. 根据节点依赖关系，即可构建出包含并行分支的树状结构

示例代码 #

import torch
import torch.fx
from your_module import CoolCNN

model = CoolCNN()
# 仅通过前向传播追踪生成计算图
traced_model = torch.fx.symbolic_trace(model)

# 遍历节点并打印依赖关系
for node in traced_model.graph.nodes:
    print(f"操作节点: {node.name} | 类型: {node.target}")
    # 提取输入依赖的节点名称
    input_nodes = [arg.name for arg in node.args if isinstance(arg, torch.fx.Node)]
    if input_nodes:
        print(f"  依赖输入节点: {input_nodes}")
    print("---")

通过该方法，main_path和parallel_path会被识别为两条独立的节点链，它们的输出会作为后续平均操作的输入，完美还原你需要的树状分支结构。

方案2：自定义张量追踪工具 #

如果需要更定制化的结构输出，可以手动实现张量依赖追踪：

1. 自定义继承torch.Tensor的包装类，添加source（生成操作）和inputs（输入张量）属性
2. 重载模型中用到的所有张量操作（ReLU、Conv2d、池化等），返回带追踪标记的张量
3. 从输出张量递归回溯inputs属性，构建完整的树状计算结构

简化示例思路 #

class TracedTensor(torch.Tensor):
    @staticmethod
    def __new__(cls, data, source=None, inputs=None):
        tensor = super().__new__(cls, data.shape, dtype=data.dtype, device=data.device)
        tensor.data = data.data
        tensor.source = source  # 记录生成该张量的操作名称
        tensor.inputs = inputs or []  # 记录输入的TracedTensor列表
        return tensor

    # 重载常用操作，保持追踪链
    def relu(self):
        return TracedTensor(torch.relu(self), source="ReLU", inputs=[self])
    
    def max_pool2d(self, kernel_size, stride):
        return TracedTensor(torch.max_pool2d(self, kernel_size, stride), source="MaxPool2d", inputs=[self])
    
    # 按需重载加法、卷积、线性层等操作...

修改模型前向传播逻辑，使用这些重载后的操作，最终从输出张量回溯所有依赖即可生成树状结构。

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