How to know input/output layer names and sizes for Pytorch model?

Question

I have Pytorch model.pth using Detectron2's COCO Object Detection Baselines pretrained model R50-FPN. I am trying to convert the .pth model to onnx.

My code is as follows.

import io
import numpy as np

from torch import nn
import torch.utils.model_zoo as model_zoo
import torch.onnx
from torchvision import models

model = torch.load('output_object_detection/model_final.pth')
x = torch.randn(1, 3, 1080, 1920, requires_grad=True)#0, in_cha, in_h, in_w
torch_out = torch_model(x)
print(model)
torch.onnx.export(torch_model,               # model being run
                  x,                         # model input (or a tuple for multiple inputs)
                  "super_resolution.onnx",   # where to save the model (can be a file or file-like object)
                  export_params=True,        # store the trained parameter weights inside the model file
                  opset_version=10,          # the ONNX version to export the model to
                  do_constant_folding=True,  # whether to execute constant folding for optimization
                  input_names = ['input'],   # the model's input names
                  output_names = ['cls_score','bbox_pred'], # the model's output names
                  dynamic_axes={'input' : {0 : 'batch_size'},    # variable lenght axes
                                'output' : {0 : 'batch_size'}})

Is it correct way to convert ONNX model? If it is the right way, how to know input_names and output_names?

Used netron to see input and output, but the graph doesn't show input/output layers.

End genocide - save Gaza · Accepted Answer · 2021-03-13 10:58:26Z

6

Try this:

import io
import numpy as np
from torch import nn
import torch.utils.model_zoo as model_zoo
import torch.onnx
from torchvision import models    

model = torch.load('model_final.pth')
model.eval()
print('Finished loading model!')
print(model)
device = torch.device("cpu" if args.cpu else "cuda")
model = model.to(device)

# ------------------------ export -----------------------------
output_onnx = 'super_resolution.onnx'
print("==> Exporting model to ONNX format at '{}'".format(output_onnx))
input_names = ["input0"]
output_names = ["output0","output1"]
inputs = torch.randn(1, 3, 1080, 1920).to(device)

torch_out = torch.onnx._export(model, inputs, output_onnx, export_params=True, verbose=False,
                               input_names=input_names, output_names=output_names)

edited Mar 13, 2021 at 10:58

End genocide - save Gaza

25k10 gold badges113 silver badges133 bronze badges

answered Nov 1, 2020 at 6:01

Prajot Kuvalekar

6,7943 gold badges28 silver badges34 bronze badges

Sign up to request clarification or add additional context in comments.

1 Comment

user3303020 Over a year ago

How do you decide the input_names and output_names?

Collectives™ on Stack Overflow

How to know input/output layer names and sizes for Pytorch model?

1 Answer 1

1 Comment

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

1 Comment

Your Answer

Sign up or log in

Post as a guest

Related