Failing to export a mask rcnn model to .onnx format

[Kodemannen]

I have a trained Mask R-CNN model that I am trying to export to .onnx format, but I am having some trouble with it.

One loose guess I have is that the error stems from the model expecting dictionaries as input, which torch.onnx.export() may not like, but I don't really know.

The code is as follows:

import numpy as np
import matplotlib.pyplot as plt
import torch

from detectron2.modeling import build_model
from detectron2.checkpoint import DetectionCheckpointer

from openslide import OpenSlide

import os
from detectron2.config import get_cfg

device = torch.device("cpu")

def add_new_fields_to_config(cfg):
    cfg.INPUT.BRIGHTNESS_DELTA = 0.0
    cfg.INPUT.CONTRAST_DELTA = 0.0
    cfg.INPUT.SATURATION_DELTA = 0.0
    cfg.INPUT.HUE_DELTA = 0.0
    cfg.INPUT.HE_NORM = False
    cfg.DATASETS.NUM_TRAIN_IMAGES = 1000    # Just some number to be overwritten by config 
    cfg.DATASETS.TEST_JSON = ("",)
    cfg.SOLVER.NUM_EPOCHS = 1
    cfg.SOLVER.STEP_SIZE_EPOCHS = 10  # Reduce learning rate by gamma every steps size epochs
    return cfg


cfg = get_cfg()
cfg = add_new_fields_to_config(cfg)
yaml_config_path = os.path.expanduser("~/data/detectron2-data/model/config.yaml")
cfg.merge_from_file(yaml_config_path)

model = build_model(cfg).to(device)  

weights_path = os.path.expanduser("~/data/detectron2-data/model/training_5_checkpoint_8414.pth")
DetectionCheckpointer(model).load(weights_path)  

model.eval()
batch_size = 1


# Create some artificial input data for tracing:
data = [torch.randn(3, 224, 224, requires_grad=False)]*2
inputs =  [{"image": x} for x in data]

torch_out = model(inputs)


torch.onnx.export(model,
                  inputs,
                  "training_5_checkpoint_8414.onnx",
                  export_params=True,
                  opset_version=10,
                  do_constant_folding=True,
                  input_names=['input'],
                  output_names=['output'],
                  dynamic_axes={'input' : {0 : 'batch_size'},
                                'output' : {0 : 'batch_size'}}
                  )

The error I get is this:

Traceback (most recent call last):
  File "/home/samuel/repos/InstanceSegmentation/detectron2_implementation/export-to-onnx/export-to-onnx.py", line 77, in <module>
    torch.onnx.export(model,
  File "/home/samuel/.local/lib/python3.10/site-packages/torch/onnx/__init__.py", line 350, in export
    return utils.export(
  File "/home/samuel/.local/lib/python3.10/site-packages/torch/onnx/utils.py", line 163, in export
    _export(
  File "/home/samuel/.local/lib/python3.10/site-packages/torch/onnx/utils.py", line 1074, in _export
    graph, params_dict, torch_out = _model_to_graph(
  File "/home/samuel/.local/lib/python3.10/site-packages/torch/onnx/utils.py", line 727, in _model_to_graph
    graph, params, torch_out, module = _create_jit_graph(model, args)
  File "/home/samuel/.local/lib/python3.10/site-packages/torch/onnx/utils.py", line 602, in _create_jit_graph
    graph, torch_out = _trace_and_get_graph_from_model(model, args)
  File "/home/samuel/.local/lib/python3.10/site-packages/torch/onnx/utils.py", line 517, in _trace_and_get_graph_from_model
    trace_graph, torch_out, inputs_states = torch.jit._get_trace_graph(
  File "/home/samuel/.local/lib/python3.10/site-packages/torch/jit/_trace.py", line 1182, in _get_trace_graph
    outs = ONNXTracedModule(f, strict, _force_outplace, return_inputs, _return_inputs_states)(*args, **kwargs)
  File "/home/samuel/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/samuel/.local/lib/python3.10/site-packages/torch/jit/_trace.py", line 134, in forward
    graph, out = torch._C._create_graph_by_tracing(
  File "/home/samuel/.local/lib/python3.10/site-packages/torch/jit/_trace.py", line 124, in wrapper
    out_vars, _ = _flatten(outs)
RuntimeError: Only tuples, lists and Variables are supported as JIT inputs/outputs. Dictionaries and strings are also accepted, but their usage is not recommended. Here, received an input of unsupported type
: Instances

Edit: Issue #3224 seems to be about the same error. Should I have made this post as a comment in that thread instead?