Exporting MMDetection models to ONNX format
MMDetection is an open-source object detection toolbox based on PyTorch. This article explains how to export MMDetection models to ONNX format for use with the ailia SDK.
Structure of the MMDetection model
The MMDetection model consists of two files: a
config file that represents the design of the neural network, and a
checkpoint file that represents the trained parameters.
The config file is a text file written in Python, where variable names and data structures are set according to the MMDetection definition rules. Here is an example of such config file.
The structure of the config file is explained in the MMDetection documentation site.
Config System — MMDetection 2.4.0 documentation
To help the users have a basic idea of a complete config and the modules in a modern detection system, we make brief…
Official conversion scripts
There is an official conversion script available, which can be used to export MMDetection models to ONNX format.
python3 tools/pytorch2onnx.py <config file> <checkpoint file> --out <out.onnx> --shape 1120 768
Specify the paths of the config and checkpoint files as arguments, the name of the output ONNX file with the
--out parameter, and the shape of the input tensor with the
MMDetection does not yet have sufficient support for ONNX export, therefore exporting using
pytorch2onnx.py may not work.
We will use
OTEDetection, which was developed based on MMDetection to support ONNX export of many models, including SSD, FCOS, ATSS, FoveaBox, Faster & Mask R-CNN, Cascade & Cascade Mask R-CNN.
OTEDetection — openvinotoolkit/mmdetection
This is an Object Detection and Instance Segmentation toolbox, that is a part of OpenVINO Training Extensions. Project…
End-to-end Faster/Mask R-CNN models export to ONNX by druzhkov-paul · Pull Request #1386 ·…
Since primitives like ROIAlign and NonMaxSuppression required for most of the detection/instance segmentation models…
The OTEDetection conversion script runs as follows.
python3 tools/export.py <config file> <checkpoint file> <output dir> onnx
Update to Version 2 format
OTEDetection is based on version 2 of MMDetection and is not compatible with version 1 format model files. In order to export a model file created in version 1 format to ONNX format, it must first be converted to version 2 format.
The model’s config file can be edited in a text-based format by referring to the instructions on the MMDetection documentation site; a detailed description of the config file can be found here.
For the checkpoint file, there is a conversion tool available, which can be used to run the following.
python tools/upgrade_model_version.py <checkpoint file v1> <output file>