YOLOv5モデルをONNXモデルにして使いたいけど後処理が面倒なとき

困ったこと

YOLOv5は便利なライブラリですが、ONNXへモデルを変換したときにちょっと困ったことがあります。
というのも、変換後のONNXモデルにはNMSなどの後処理が含まれていないため、後処理は別途用意する必要があります。
公式ではPyTorchの関数を使ったNMSになっているため、そのまま後処理のコードをコピーしようとすれば実行環境上にONNX RuntimeとPyTorchの両方を用意しないといけません。でもせっかくONNXを使うなら、環境にPyTorchを入れたくないですよね。

解決方法

PyTorchを入れたくないけどどうしよう…と困っていたところ、こちらのプルリクを見つけました。
https://github.com/ultralytics/yolov5/pull/7736
どうやらNMSの処理がONNXモデルに含まれるような修正をおこなっているようです。

2022/07/17現在はまだmasterへはマージされていないのですが、fork先のブランチを試してみると、うまくいくことが確認できました。
実際にONNXモデルへ変換をおこなうときにはexport.pyに"–nms"オプションをつければOKです。 モデルの出力値の扱いはこちらを参考にすると分かるかと思います。

comments powered by Disqus
Built with Hugo
テーマ StackJimmy によって設計されています。