finding the input size for detectron2 model to convert to onnx

Question

I am trying to convert this detectron 2 model to onnx format but facing issues with the dummy input shape. here is my code:

import torch
import cv2
from detectron2.config import get_cfg
from detectron2.modeling import build_model
from detectron2.checkpoint import DetectionCheckpointer
from detectron2 import model_zoo

cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
cfg.MODEL.DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

model = build_model(cfg)
model.eval()
model.to(cfg.MODEL.DEVICE)

DetectionCheckpointer(model).load(cfg.MODEL.WEIGHTS)

dummy_input = torch.randn(1, 3, 224, 224, device=cfg.MODEL.DEVICE)

print('input: ',dummy_input.shape)

input_names = ["actual_input_1"]
output_names = ["output1"]

torch.onnx.export(model, dummy_input, "detectron2_model.onnx", verbose=True,
                  input_names=input_names, output_names=output_names)

print("Detectron2 model has been successfully exported to detectron2_model.onnx")

Error:

input:  torch.Size([1, 3, 224, 224])
---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
<ipython-input-10-569a1c2752b7> in <cell line: 33>()
     31 output_names = ["output1"]
     32 
---> 33 torch.onnx.export(model, dummy_input, "detectron2_model.onnx", verbose=True,
     34                   input_names=input_names,
     35                   output_names=output_names)

16 frames
/content/detectron2/detectron2/modeling/meta_arch/rcnn.py in <listcomp>(.0)
    225         Normalize, pad and batch the input images.
    226         """
--> 227         images = [self._move_to_current_device(x["image"]) for x in batched_inputs]
    228         images = [(x - self.pixel_mean) / self.pixel_std for x in images]
    229         images = ImageList.from_tensors(

IndexError: too many indices for tensor of dimension 3

I want the code to work and find the optimal dummy input shape to convert to onnx format. Your help is appreaciated. Thanks in advance.

Answer 1

Have you checked the input shape of detectron2_model.onnx?

You can try and adjust this script to find the shape of the .onnx model and match your input against it,

import onnx
onnx_model = onnx.load("detectron2_model.onnx")
for i in onnx_model.graph.input:
    print(i.name)
i_shape = onnx_model.graph.input[0].type.tensor_type.shape.dim
print([d.dim_value for d in i_shape])

References:

1. https://onnx.ai/onnx/intro/python.html
2. Find input shape from onnx file