facebookresearch / facebookresearch/detectron2

Export detectron2 models to ONNX retaining the batch-norm layers.

Open
#3,229 0 comments 0 reactions 0 assignees View on GitHub
enhancement
Dominant language
Python
Stars
34.7k
Forks
7.9k
PR merge metrics
No merged PRs in 30d

Description

## 🚀 Feature
Option to export detectron2 models using onnx.export.Caffe2Tracer and export_onnx() while retaining the batch-norm layers.

## Motivation & Examples

Currently, the following export code for detectron2's maskrcnn to onnx, optimizes out the batch-norm layer-

```
modelName = "R_50_C4_3x"
modelYaml = "COCO-InstanceSegmentation/mask_rcnn_"+modelName+".yaml"
torch_model = model_zoo.get(modelYaml, trained=True)

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, 'maskrcnn.onnx')

```
My use case involves fine-tuning the mask-rcnn model in a different framework, for which I would need the explicit non-fused normalization layer.

Given the Caffe2 tracer maintains the frozen batch-norm operation on the trace, it'll be extremely useful to have an option on export_onnx() , to turn off conv-batchnorm fusion/optimization.

Is there a way to accomplish this using the existing framework?

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.