如何从PyTorch的.pth模型中提取含激活函数的完整网络结构

当激活函数以F.relu这类functional API形式写在forward方法中时，它们不会被视为模型的子模块，因此直接打印模型结构无法看到。要提取这些操作，需通过追踪模型的前向传播过程捕获所有计算步骤，以下是具体实现方法：

方法1：使用TorchScript追踪模型 #

TorchScript可将模型的前向传播转化为可序列化的中间表示，能完整记录所有调用的操作，包括functional形式的激活函数。

步骤： #

1. 加载保存的模型文件

   import torch
   import torch.nn.functional as F
   
   # 加载完整模型（基于torch.save(model)保存的文件）
   model = torch.load("your_model.pth")
   model.eval()  # 切换为评估模式
   

   注意：若加载时提示模型类未定义，说明.pth仅保存了state_dict，需先恢复模型结构（无源码时可参考方法2补充）。

2. 创建与模型输入维度匹配的张量作为追踪输入

   # 示例：假设模型输入为3通道、256x256的图像，需根据实际场景调整
   dummy_input = torch.randn(1, 3, 256, 256)
   

3. 追踪模型并生成TorchScript模块

   traced_model = torch.jit.trace(model, dummy_input)
   

4. 查看完整计算图

   • 直接打印追踪后的模型操作细节：

     print(traced_model.graph)
     

   • 保存为TorchScript文件后可视化（更直观）：

     traced_model.save("traced_model.pt")
     # 可用Netron工具打开.pt文件，清晰查看所有层与操作节点
     

   输出的graph中会明确显示aten::relu、aten::flatten等所有在forward中调用的操作。

方法2：导出为ONNX格式并可视化 #

ONNX格式会完整记录模型的计算图结构，包含所有functional操作，适合无源码场景下的模型分析。

步骤： #

1. 加载模型并设置为评估模式（同方法1）
2. 导出模型为ONNX文件

   dummy_input = torch.randn(1, 3, 256, 256)  # 匹配模型实际输入维度
   torch.onnx.export(
       model,
       dummy_input,
       "model.onnx",
       opset_version=13,  # 选择适配的opset版本
       do_constant_folding=True,
       input_names=["input"],
       output_names=["output"],
       verbose=True  # 打印导出的操作细节
   )
   

3. 可视化ONNX模型
   使用Netron工具（本地客户端或在线版本）打开model.onnx，可直观查看每一层的操作，包括Relu、Flatten等在forward中调用的函数。

注意事项 #

• 确保dummy_input的维度与模型实际输入完全匹配，否则追踪/导出会失败。若未知输入维度，可从模型第一层参数反推（比如卷积层的in_channels）。
• 若.pth仅保存了state_dict（而非完整模型），可尝试从state_dict的键名推断模型结构（如conv1.weight对应卷积层），或结合torchinfo工具尝试重建模型框架后再进行追踪。

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