Export detectron2 pre-trained mask-rcnn to ONNX with BatchNorm.

coderAD: <p>Hi folks,</p>
<p><strong>BLOT</strong>: Need help exporting detectron2’s maskrcnn to ONNX along with the frozen batch norm layers.</p>
<p>I’m fairly new to detectron2 framework and had some issues exporting detectron2’s mask-rcnn to onnx, retaining the frozen batch norm layers from the torch model.</p>
<p>I have been successful in importing the resnet-50 mask-rcnn network using the code snippet below. But in this case, the frozen batch norm layers get optimized out/ constant-folded in the exported ONNX network.</p>
<pre><code class="lang-auto">modelName = "R_50_C4_3x"
modelYaml = "COCO-InstanceSegmentation/mask_rcnn_"+modelName+".yaml"
torch_model = model_zoo.get(modelYaml, trained=True)
cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file(modelYaml))
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(modelYaml)

# COCO dataset loader
data_loader = build_detection_test_loader(cfg, cfg.DATASETS.TEST[0], num_workers=0)
first_batch = next(iter(data_loader))

tracer = Caffe2Tracer(cfg, torch_model, first_batch)
onnx_model = tracer.export_onnx()
onnx.save(onnx_model, "mrcnn_onnx.onnx"))
</code></pre>
<p>Based off a conversation from <a href="https://github.com/pytorch/pytorch/issues/49226" rel="noopener nofollow ugc">this thread</a>, my understanding is that, <code>using the torch.onnx.export()</code> function, we can turn off constant_folding and optimizations, when exported using the training=TrainingMode.TRAINING option. After which, the frozen batch norm layer wouldn’t get optimized out. However, while trying this out, I run into an error (pasted below)</p>
<p>Updated export code:</p>
<pre><code class="lang-auto">torch.onnx.export(torch_model, first_batch, "mrcnn.onnx", verbose=True, do_constant_folding=True) #, training=TrainingMode.TRAINING)
</code></pre>
<p>Error message:</p>
<pre><code class="lang-auto">File "/Users/ad/anaconda3/envs/torch/lib/python3.9/site-packages/torch/onnx/utils.py", line 632, in _export
    _model_to_graph(model, args, verbose, input_names,
  File "/Users/ad/anaconda3/envs/torch/lib/python3.9/site-packages/torch/onnx/utils.py", line 409, in _model_to_graph
    graph, params, torch_out = _create_jit_graph(model, args,
  File "/Users/ad/anaconda3/envs/torch/lib/python3.9/site-packages/torch/onnx/utils.py", line 379, in _create_jit_graph
    graph, torch_out = _trace_and_get_graph_from_model(model, args)
  File "/Users/ad/anaconda3/envs/torch/lib/python3.9/site-packages/torch/onnx/utils.py", line 342, in _trace_and_get_graph_from_model
    torch.jit._get_trace_graph(model, args, strict=False, _force_outplace=False, _return_inputs_states=True)
  File "/Users/ad/anaconda3/envs/torch/lib/python3.9/site-packages/torch/jit/_trace.py", line 1149, in _get_trace_graph
    outs = ONNXTracedModule(f, strict, _force_outplace, return_inputs, _return_inputs_states)(*args, **kwargs)
  File "/Users/ad/anaconda3/envs/torch/lib/python3.9/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/Users/ad/anaconda3/envs/torch/lib/python3.9/site-packages/torch/jit/_trace.py", line 126, in forward
    graph, out = torch._C._create_graph_by_tracing(
  File "/Users/ad/anaconda3/envs/torch/lib/python3.9/site-packages/torch/jit/_trace.py", line 120, in wrapper
    out_vars, _ = _flatten(outs)

RuntimeError: Only tuples, lists and Variables supported as JIT inputs/outputs. Dictionaries and strings are also accepted but their usage is not recommended. But got unsupported type Instances
</code></pre>
<p>I faced a similar issue with <code>_flatten(in)</code> , which was resolved by converting all the scalars to tensors. But here, the outs from detectron2 is an object of type ‘Instances’.  A few questions here -</p>
<ol>
<li>Why does the exporter care about the final format after post-processing?</li>
<li>Any suggestions on how can I resolve this error? or,</li>
<li>Any leads on how to export detectron2’s maskrcnn to ONNX with the batch norm layers will be much appreciated.</li>
</ol>
<p>Thanks in advance!</p>
<p>This <a href="https://github.com/pytorch/pytorch/issues/52867" rel="noopener nofollow ugc">issue</a> is addressing the export of fixed batch norm layers from the model. But, currently I’m unable to successfully export any form of the network using the<code> torch.onnx.export(torch_model, first_batch, "mrcnn.onnx")</code> syntax.</p>