Python 中的多个 ONNX 输出（获取 ONNX 的中间层输出）

Question

如何将模型导出到 ONNX，以便获得中间层的输出以及层？ （我见过一个类似的问题，但没有得到解答这里）

假设我有一个模型，

model

。该模型是一个火炬模型，我希望有多个输出：最后一层以及中间层之一：具体来说，是过程中发生的卷积之一。

    import torch
    import onnx

    device = 'cpu'
    dummy_input = torch.randn(1, 3, 320, 320).to(device)
    input_key_name = 'input'
    output_key_name = 'output'
    torch.onnx.export(model, dummy_input, model_output_name, 
                      input_names=[input_key_name], output_names=[output_key_name])

我的问题是：

1. 是否可以得到多层输出？怎么办？
2. 我如何知道我应该提供的输出层名称？是否可以使用 Netron 来解释这些名称？ （或者其他工具？）

现在我的代码在最后一层可以正常工作，但我不确定如何从这里开始获得附加层。

Answer 1

有几种方法可以做到这一点。可以将一个模型划分为多个子模型。或者，可以修改 ONNX 图并导出具有多个输出的模型：

import torch
import torchvision
import onnx
from onnx import helper

# load pretrained ResNet model
resnet = torchvision.models.resnet50(weights="ResNet50_Weights.IMAGENET1K_V2")

# export torch model to ONNX
tensor_x = torch.rand((1, 3, 224, 224), dtype=torch.float32)
onnx_program = torch.onnx.dynamo_export(resnet, tensor_x)

model = onnx_program.model_proto
# or load an existing ONNX model
# model = onnx.load("resnet50_imagenet.onnx")

# Add an additional output
intermediate_layer_value_info = helper.make_tensor_value_info(
    name='layer4_1',  # This should be exactly equal to your intermediate tensor name in the ONNX graph
    elem_type=onnx.TensorProto.FLOAT, 
    shape=[1,2048,7,7]  # This should be the shape of the intermediate tensor
)

# Append the additional output to the model
model.graph.output.append(intermediate_layer_value_info)

# Save the Model
onnx.save(model, "resnet50_imagenet_modified.onnx")

目标层的名称可以通过打印火炬模型获得，也可以通过使用Python打印模型或上传到netron

从ONNX图中获得