MMDetectionのモデルをONNX形式にエクスポートする

MMDetectionは、PyTorchに基づくオープンソースのオブジェクト検出ツールボックスです。MMDetectionのモデルをailia SDKで使用するONNX形式にエクスポートする手順について解説します。ailia SDKについてはこちらをご覧ください。

出典：https://github.com/open-mmlab/mmdetection

MMDetectionモデルの構造

MMDetectionのモデルは、ニューラルネットワークの設計を現す ”configファイル” と学習済みのパラメータを現す ”checkpointファイル” の二つのファイルから構成されます。

configファイルはPythonのコード形式で記述されたテキストファイルで、変数名やデータ構造をMMDetectionの定義ルールに従って記述します。例えば、以下のような内容になります。

configファイルの構造は、MMDetectionのドキュメントサイトに解説があります。（ドキュメントサイトのリンクはこちら）

Config System - MMDetection 2.4.0 documentation

公式の変換スクリプト

公式の変換スクリプトが用意されていて、これを使ってMMDetectionのモデルをONNX形式にエクスポートすることができます。（スクリプトのリンクはこちら）

python3 tools/pytorch2onnx.py <config file> <checkpoint file> --out <out.onnx> --shape 1120 768

引数に ”configファイル” と ”checkpointファイル” のパスを指定し、 --out パラメータで作成するONNXファイルの名前を指定し、 --shape パラメータで入力テンソルのshapeを指定します。

OTEDetection

MMDetectionは、ONNXのエクスポートに対してまだ充分に対応していないため、 pytorch2onnx.pyを使ったエクスポートはうまくいかないかもしれません。

そこでMMDetectionをベースに開発され、多くのモデルのONNXエクスポートに対応できるようにした OTEDetection を使用します。OTEDetectionは、SSD, FCOS, ATSS, FoveaBox, Faster & Mask R-CNN, Cascade & Cascade Mask R-CNN など多くのモデルのONNXエクスポートに対応しています。

OTEDetection - openvinotoolkit/mmdetection

End-to-end Faster/Mask R-CNN models export to ONNX by druzhkov-paul · Pull Request #1386 ·…

OTEDetectionの変換スクリプトは以下のように実行します。（スクリプトのリンクはこちら）

python3 tools/export.py <config file> <checkpoint file> <output dir> onnx

Version2形式へのアップデート

OTEDetectionは、MMDetectionのバージョン2をベースにしており、バージョン1形式のモデルファイルとは互換性がありません。バージョン1形式で作成されたモデルファイルをONNX形式にエクスポートするためには、まずバージョン2形式へ変換を行う必要があります。

モデルのconfigファイルは、MMDetectionのドキュメントサイトの説明を参考に、テキストベースでの編集を行います。configファイルの詳細な説明はこちらをご覧ください。

checkpointファイルは、変換ツールが用意されているので、それを使って以下のように実行します。（変換ツールのリンクはこちら）

python tools/upgrade_model_version.py <checkpoint file v1> <output file>

ax株式会社はAIを実用化する会社として、クロスプラットフォームでGPUを使用した高速な推論を行うことができるailia SDKを開発しています。ax株式会社ではコンサルティングからモデル作成、SDKの提供、AIを利用したアプリ・システム開発、サポートまで、 AIに関するトータルソリューションを提供していますのでお気軽にお問い合わせください。