Last updated on

MMDetectionで物体検出を独自モデルで学習する! という試行錯誤 その3!


前回に引き続き、飽きずにごそごそと物体検出を試しています。

AIの推論だけに特化したONNX Runtimeを試す。

物体検出の学習にPytorch、MMDetectionを利用しているのですが、一般にこのようなライブラリは学習に重きを置いていて、ライブラリも大きなものになりがちなのですが、ONNXは推論だけに特化したモデルとのこと。
確かにONNX形式に変換、みたいな文字列も、MMDetectionを使う上で参考とさせていただいていたサイトに記載されていて、推論だけに特化していれば軽量、軽快なのかもしれない?と思い、いろいろ試していました。

 

MMDetectionで学習したデータをONNXに変換する!

MMDetectionで学習を行った重みのデータは、MMDeployというライブラリで様々なモデルに変換ができます。
が、あまり情報がなかったり、情報が古かったり、私のONNXの理解が乏しかったりで結構試行錯誤しました。
変換するだけなのになー

変換だけのためにライブラリをインストールするのも嫌だったので、Dockerでの対応にしました。
わたしがDockerをよくわかってないので、ほぼ書いてある通りにやってます。
https://mmdeploy.readthedocs.io/en/latest/01-how-to-build/build_from_docker.html

2024/08/05時点で、大まかな手順は、

  1. Dockerでイメージをpullする

    docker pull openmmlab/mmdeploy:ubuntu20.04-cuda11.8-mmdeploy1.3.1

  2. Dockerでコンテナを起動する
    docker run --gpus=all -it --rm openmmlab/mmdeploy:ubuntu20.04-cuda11.8-mmdeploy1.3.1
  3. コンテナにMMDetectionをセットアップする(滅多に変換しない予定だったので、毎回セットアップしているというアホさ)
    mim install mmdet
  4. Dockerのコンテナに、MMDetectionで学習した重み omomi.pth、学習で生成された config.py、学習に利用した画像 image.pngをコピー
    ローカルPC参照できるなら直接指定してもいいのだと思いますが、よくわかんなかったのでコピーしてます。
  5. MMDeployでONNXに変換する
    python3 mmdeploy/tools/deploy.py mmdeploy/configs/mmdet/detection/detection_onnxruntime_static.py config.py omomi.pth image.png
  6. work-dirというディレクトリにONNXに変換された重み、end2end.onnx が完成! ぱー!

deploy.pyでも、恐らくtorch2onnx.pyに行くようで、torch2onnx.pyを実行しても同じ引数を要求されるのですが、公式ではdeploy.pyと書いてあるので、書いてある通り実行しています。
https://mmdeploy.readthedocs.io/en/latest/02-how-to-run/convert_model.html

なお、変換できるリストもあります。
https://mmdeploy.readthedocs.io/en/latest/03-benchmark/supported_models.html

最初、変換でできたONNXを実行してみたら何かエラーが出て、変換できてねーじゃねーかー!DETRは未対応かー??って思ったのですが、DETR系は2023/12月に対応していました。
(入力画像の解像度が数ピクセル違ってて、エラーがでていたのでした。素人なのにプロを疑ってすみません・・・)

 

ONNXで物体検出する!

ONNXのコードをゼロから書くほどの能力はないので、下記サイトを参考にしました。

分類ラベルのテキストの周りの枠塗りがうまく動かなかったので、コメントアウトしました。
で、ONNX Runtimeを使う上で大事なのは、

image = image.resize((640, 480)) # モデルの入力サイズにリサイズ

学習に用いた画像の解像度の縦・横それぞれ1/2にした画像がONNXへのインプットになっているところ!
私がONNXをよくわかっていなかったのはここで、解像度が1/4になってしまうことで、画像上で小さいものは更に潰れて見にくくなってしまうんですよねー
私が検出したかったものは、画像の中では結構小さい部分だったので、認識率がガクっと落ちてしまいました。
試しにリサイズせずに画像をインプットしてみると、サイズが違うよってエラーが出るので、だめなんだなーと。

1/4になるなら、ONNXに変換する時に、画像を4倍にしておけばいいんじゃないか??と思いましたが、こちらも(たぶん)うまくいきませんでした。
たぶん、と書いたのは、mmdeployでdeploy.pyする時に、ほんとにちゃんと4倍にした画像を指定したよな??というのが微妙に自信がなくなったためです。(何度も、Docker立ち上げて変換して、試して、Dockerまた立ち上げて・・・を繰り返してたので、よくわからなくなってしまった)

 

ONNXを使ってみた感想

入力画像を固定でき、検出したい物体が画像を1/4にしても大きく形状などがわかる場合、ONNXはとても良いだろうと感じました。
いわゆるエッジコンピューティングで、ウェブカメラのようなものの画像を小さいパソコンで処理する、というような場合、大活躍すると感じました。
ONNXの処理はCPUで実施したのですが、かなりレスポンスよく処理されてきました。CUDA対応も可能なようなので、相当なレスポンスが可能と思います。ONNXでやりたかったなー

入力解像度を1/4にしてしまうことで、小さいものがつぶれてしまい認識率が下がってしまうのは、私にとってはとても残念な結果でした。
また、入力解像度が固定されるため、いろんな解像度に対応しようとするとONNXモデルが複数必要なことも残念ポイントでした。その程度の差、AIの中で良い感じに対応してくれてもいいのになー、みたいな。
でも推論特化なのでこういう形になるんでしょうね。
こちらは解像度ごとにアノテーションを作って、解像度ごとにモデルを変換したら、なんとかなるかな?と思いましたが、縮小での潰れによる認識率低下が致命的でした。処理は早くていいんだけどなー
ざんねん!