axinc-ai
Published in

axinc-ai

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.

Official conversion scripts

There is an official conversion script available, which can be used to export MMDetection models to ONNX format.

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 --shape parameter.

OTEDetection

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.

The OTEDetection conversion script runs as follows.

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.

ax Inc. has developed ailia SDK, which enables cross-platform, GPU-based rapid inference.

ax Inc. provides a wide range of services from consulting and model creation, to the development of AI-based applications and SDKs. Feel free to contact us for any inquiry.

--

--

Get the Medium app

A button that says 'Download on the App Store', and if clicked it will lead you to the iOS App store
A button that says 'Get it on, Google Play', and if clicked it will lead you to the Google Play store