如何知道 Pytorch 模型的 input/output 层名称和大小？

Question

我有 Pytorch model.pth 使用 Detectron2's COCO 对象检测基线预训练模型 R50-FPN。我正在尝试转换 .pth model to onnx.

我的代码如下

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'}})

转换ONNX模型的方法正确吗？如果是正确的方法，如何知道input_names和output_names？

使用 netron 查看输入和输出，但图表未显示 input/output 层。

Answer 1

试试这个：

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)

如何知道 Pytorch 模型的 input/output 层名称和大小？

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

python

pytorch

onnx

onnxruntime