如何知道Pytorch模型的输入/输出层名称和大小？

问题描述

我使用Detectron2's COCO对象检测基线预训练的模型R50-FPN具有Pytorch model.pth。我正在尝试转换.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模型是否正确？如果是正确的方法，如何知道输入名称和输出名称？

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

解决方法

暂无找到可以解决该程序问题的有效方法，小编努力寻找整理中！

如果你已经找到好的解决方法，欢迎将解决方案带上本链接一起发送给小编。

小编邮箱:dio#foxmail.com (将#修改为@）

onnx onnxruntime python pytorch