PyTorch模型导出ONNX：如何获取中间层及最终层多输出

1. 能否获取多层输出？怎么实现？ #

完全可以实现导出ONNX时同时输出多层结果，核心思路是让模型在forward过程中返回你需要的最终层输出和目标中间层输出，具体有两种常用方式：

方式一：修改原模型的forward方法 #

如果有权限修改原模型代码，直接在forward函数末尾把需要的中间层张量和最终输出一起返回：

class YourModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = torch.nn.Conv2d(3, 64, kernel_size=3)
        # 其他层定义...
        self.final_layer = torch.nn.Linear(...)
    
    def forward(self, x):
        x = self.conv1(x)
        # 中间层计算过程...
        mid_output = x  # 假设这里是你需要的中间卷积层输出
        final_output = self.final_layer(x)
        return final_output, mid_output  # 返回两个输出

方式二：创建模型包装类（无需修改原模型） #

如果不想改动原模型代码，可以写一个简单的包装类，把原模型包裹起来，在forward里提取中间层输出：

class ModelWrapper(torch.nn.Module):
    def __init__(self, original_model):
        super().__init__()
        self.model = original_model
        # 提前获取目标中间层，比如假设目标层是model.backbone.conv3
        self.target_mid_layer = original_model.backbone.conv3
    
    def forward(self, x):
        # 遍历模型层，执行前向并提取目标中间层输出
        for module in self.model.children():
            x = module(x)
            if module is self.target_mid_layer:
                mid_output = x.clone()  # 保存中间层输出
        final_output = x
        return final_output, mid_output

之后导出ONNX时，只需把output_names参数改成包含多个输出名的列表即可：

import torch
import onnx

device = 'cpu'
dummy_input = torch.randn(1, 3, 320, 320).to(device)
input_key_name = 'input'
# 定义多个输出的名称
output_key_names = ['final_output', 'mid_conv_output']

# 如果用包装类，先实例化包装后的模型
wrapped_model = ModelWrapper(model).to(device)
# 导出时指定多个输出名
torch.onnx.export(wrapped_model, dummy_input, "multi_output_model.onnx", 
                  input_names=[input_key_name], output_names=output_key_names)

2. 如何确定目标中间层的名称？ #

可以通过两种方式快速定位：

方法一：直接查看PyTorch模型结构 #

在Python里打印模型的结构，就能看到各层的名称：

print(model)
# 或者查看带完整路径的层名称
for name, module in model.named_modules():
    print(name, module)

运行后会输出类似backbone.conv3 Conv2d(...)的内容，这里的backbone.conv3就是该层的名称，直接用它来定位中间层即可。

方法二：用Netron工具查看 #

先导出原来的单输出ONNX模型，然后用Netron打开这个模型文件：

• 打开Netron后，会可视化展示模型的计算图
• 找到你需要的中间卷积层，点击它就能看到层的名称、输入输出张量信息
• 记下这个名称，就可以在包装类或原模型里对应提取输出

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