facebookresearch / facebookresearch/detectron2
Export detectron2 models to ONNX retaining the batch-norm layers.
- 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
Assessment
This issue has not been assessed yet.