【物体検出】YOLOシリーズ最新版のYOLOXを試してみる 〜デモから学習まで〜

はじめに

物体検出でお馴染みのYOLOシリーズの最新版「YOLOX」が登場しましたので早速試していきます。今回は本家のチュートリアルに沿って推論デモとcocoデータセットの学習を紹介します。Google colab環境を使用しますので、誰でも簡単に試すことが可能です。

(2022年7月にYOLOv7が登場しております。詳細はリンクよりご確認ください)

YOLOXとは

YOLOXは2021年8月に公開された最新バージョンであり、速度と精度の面で限界を押し広げています。YOLOv5を超える性能で、使いやすいApache Licenseとなっています。YOLOXのベンチマーク結果は以下となっています。

YOLOXには高精度向けのStandard Modelsと、エッジデバイス用のLight Modelsの二種類が存在します。

表1　Standard Models

表2　Light Models

引用：https://github.com/Megvii-BaseDetection/YOLOX

YOLOXのアーキテクチャ

YOLOXは従来のYOLOをアンカーフリーに変更し、Decoupled HeadとSimOTAを導入した物体検出モデルです。CVPR2021の「Automous Driving WorkshopのStreaming Perception Challenge」で優勝しました。従来のYOLOv4およびYOLOv5はアンカーを使用することを前提としており、過度に最適化されたパイプラインとなってしました。一方で、YOLOXはYOLOv3-SPPをベースラインとして改良しています。また、YOLOv3-SPPに対してYOLOv5で導入された、CSPNetのバックボーンとPAN headを導入しています。
物体検出においては、regression（バウンディングボックスの位置の計算）とclassification（クラス分類）を同時に行っており、Decoupled headを使用する方法が提案されています。これまでのYOLOシリーズのbackboneとfeature pyramidsにはCoupled headが使用されてましが、YOLOXではCoupled headをDecoupled headに置き換えることで精度の向上を図っています。学習においては、MosaicやMixupによるData Augmentationを適用しています。

環境構築

早速YOLOXを使用していきましょう。以下、Google colab環境で進めていきます。
まずは本家よりcloneしてきます。

from google.colab import drive
drive.mount('/content/drive')

%cd ./drive/MyDrive
!git clone https://github.com/Megvii-BaseDetection/YOLOX

次に必要なライブラリをインポートします。

%cd YOLOX
!pip install -U pip && pip install -r requirements.txt
!pip install -v -e . 

こちらも合わせて準備します。

!pip install cython
!pip install 'git+https://github.com/cocodataset/cocoapi.git#subdirectory=PythonAPI'

以上で準備完了です。
YOLOXを使えるようになりました。

推論デモ

まずは学習済みモデルを取得します。
冒頭紹介した通り、いくつかの学習済みモデルがあります。
ここでは、Standard Modelsの中で最も高性能であるyolox_xを使うことにします。

!wget https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_x.pth

YOLOXではデモ用のコードが用意されいています。
これをそのまま実行することで簡単に物体検出を試すことができます。
画像はassets/dog.jpgを指定しています。必要によりconf、nms、tsizeの設定を変更してください。

python

!python tools/demo.py image -n yolox-x -c yolox_x.pth --path assets/dog.jpg --conf 0.25 --nms 0.45 --tsize 640 --save_result --device gpu

YOLOX_outputs/yolox_x/vis_resに結果が保存されます。
中身を確認してみると、無事に犬や自転車が検出できていることがわかります。
yoloシリーズではおなじみの画像ですね。

画像だけだなく動画でもテストすることが可能です。
よかったら動画もご覧ください。

データセットと学習済モデル

ここから学習の方法を紹介していきます。
データセットは公式で推奨されている通り「mini-coco128」のデータセットを使用します。
早速データを取得します。

取得が完了したらdatasetsに移動させます。
フォルダの中を見ると、YOLOXではcoco形式のデータセットとなっていることがわかります。

次に学習済モデルを取得します。今回はチュートリアルに従って、yolox_s.pthを使用します。先ほどと同様に学習済モデルを取得します。

!wget https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_s.pth

以上で学習に必要なデータがそろいました。

学習

いよいよ学習を開始します。以下の通り実行します。

!python tools/train.py -f exps/example/custom/yolox_s.py -d 1 -b 8 --fp16 -o -c yolox_s.pth

-d:GPUデバイスの数
-b:バッチの合計サイズ
「mini-coco128」は比較的小規模データセットのため、すぐに学習が終わります。
結果はYOLOX_outputs/yolox_s/に保存されます。

評価

このあとは作成したモデルでテストしてみましょう…となるころですが、データセットのラベル名を設定する必要があります。デモの通りに推論を進めると、今回使用したデータセットとは異なるデフォルトで設定されているラベルが表示されます。テスト結果を表示しても正しいラベルが表示されないため、（ラベルを書き換えればいいのだけですが手間なので）　ここでは割愛します。具体的な評価方法は次回の記事を参照してください。

まとめ

最後までお読みいただきありがとうございました。
YOLOX本家のチュートリアルに沿って推論デモとcocoデータセットを用いた学習を紹介しました。
今回のシリーズも推論から学習まで、非常に使いやすいものになっています。
Light Modelsの精度がさらに向上すればエッジAI（物体検出）がさらに身近なものになりそうですね。

続き

YOLOXでカスタムデータセットの学習及び推論の方法をまとめます。
実際にデータセットを作成して学習していきます。
ぜひ合わせてご覧ください。

（↓こちらではより詳細に紹介しております。よかったらご覧ください。）