YOLOXの学習をCOCOのデータセットから試してみた

本サイトはアフィリエイト広告を利用しています。
AI

YOLOXでの学習を実際に一から試してしてみました。

COCOのデータセットの入手から、物体検出まで行っています。

Colab上で実施していますので、誰でも同様にテスト可能です。

 

スポンサーリンク

YOLOXの学習をCOCOのデータセットから試してみた

YOLOX-tinyを使って、COCOのデータセットを学習させてみました。

Colab上でデータセットの入手から、学習させたモデルで画像認識まで行っています。

 

ColabのGPUを使用します

1000枚程度の写真を使ってYOLOXの学習をしますので、GPUが必要になります。

ランタイムからGPUを選択してColabを開始します。

(下記はGPUがセットされているかの確認です)

 

学習で数時間かかります。無料のColabだと運が悪いと制限が掛かる場合があります。

下記記事のように筆者は有料のColab Pro環境で動かしています。

GPUバックエンドに接続できませんが出たのでColab Proに課金した話

GPUバックエンドに接続できませんが出たのでColab Proに課金した話 
最近Google Colabで重いデータセットを学習させるケースが増えてきました。 ただし無料版のColabで長時間GPUを使っていると、制限が掛かり暫く使えなくなります。 そのため有料版のColab Proに課金してみました。

 

FiftyOneでCOCOのデータセットを入手します

COCOのデータセットはFiftyoneというツールを使って入手します。

 

Colab上で簡単にダウンロード、ラベルの整理が行えます。下記で使い方を紹介しています。

COCOのデータセットの簡単な使い方!FiftyOneを試してみた

COCOのデータセットの簡単な使い方!FiftyOneを試してみた
COCOのデータセットを簡単な使い方を紹介します。 Fifityoneというツールを使えば、任意のデータセットを簡単に作成可能です。 Colab上で動かしましたので、誰でも同様にテストできます。

 

学習時間を数時間に抑えるため、今回は1000枚程度の写真で学習させます。

(COCOのデータセットの全てだと20GB近くなるので…)

 

今回は人・車・椅子の3つを学習させるようにします。

 

YOLOX-Tinyでの学習

YOLOXの中にも何個かのモデルに分かれています。

詳細はYOLOXのGithubを見てもらえればと思います。(というか筆者が説明できない)

 https://github.com/Megvii-BaseDetection/YOLOX

 

今回は軽量モデル(Light Models)のYOLOX-tinyを使って学習していきます。

ダウンロードしたデータを学習できるように整理します。

 

YOLOXのインストールは公式通りです。

 

YOLOXフォルダにあるYOLOX-tiny.pyのファイルを修正します

 

Colab上でも修正できますので、データセットのパスを追記します。

 

あとはYOLOXの学習をさせるだけです。2~3時間ぐらいかかりました。

 

学習したモデルで画像認識します

学習済のモデルは下記に保存されました。

/content/YOLOX/YOLOX_outputs/yolox_tiny/best_ckpt.pth

学習したモデルで画像認識できるか一度テストしています。

 

下記はテスト時にOPENCV関連のエラーが出たので再インストールしているだけです。

 

実際にテスト用の画像データとモデルパスを貼り付けて、テストさせます。

 

テストすると画像認識された結果が下記フォルダに入っていました。

/content/YOLOX/YOLOX_outputs/yolox_tiny/vis_res/***/***.jpg"

***箇所に関しては日付やファイル名で適宜変わります。

 

テスト用の画像データから、人(person)をしっかり認識できました。

 

Googleドライブに学習結果を保存します

必要に応じて、出力結果をGoogleドライブに保存します。

(Colabで作成されたファイル・フォルダは閉じてしまうと保存されないため)

 

参考記事

YOLOX+Colabの学習に関して、下記記事を非常に参考にさせていただきました。

この場を借りてお礼申し上げます。

 https://github.com/Kazuhito00/YOLOX-Colaboratory-Training-Sample

 

GitHubのGist

同様にテストしたい、実行履歴を確認したい方はGitHubのGistが下記となっています。

ご自由にお使い下さい。

まとめ

YOLOX+COCOの学習を紹介させていただきました。

よろしければ皆さまも試してみて下さい。

コメント