Multiple ONNX outputs (Get Intermediate layer output for ONNX) in python

Question

How can I export model to ONNX so that I'll get intermediate layers' output as well as the layer? (I've seen a similar question that went unanswered here)

Consider I have a model, model. The model is a torch model and I'd like to have multiple outputs: the last layer as well as one of the intermediate layers: specifically, one of the convolutions that happens in the process.

    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])

My questions are:

Is it possible to get multiple layers output? How?
How would I know the output layer name I'm supposed to provide? Is it possible to use Netron to elucidate the names? (or other tools?)

Right now my code works correctly for the last layer, but I'm not sure how to go from here to get an additional layer as well.

Aray Karjauv · Accepted Answer · 2024-02-15 14:13:23Z

There are a couple of ways to do this. One can divide a model into multiple sub-models. Alternatively, one can modify ONNX graph and export a model with multiple outputs:

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")

The name of the target layer can be obtained by printing the torch model or from the ONNX graph by printing the model using Python or by uploading it to netron

Collectives™ on Stack Overflow

Multiple ONNX outputs (Get Intermediate layer output for ONNX) in python

1 Answer 1

Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related