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

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

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

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

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

 

スポンサーリンク

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

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

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

 

ColabのGPUを使用します

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

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

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

# 残り時間の表示
!cat /proc/uptime | awk '{printf("残り時間 : %.2f", 12-$1/60/60)}'

# GPUの確認
!nvidia-smi

 

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

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

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

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

 

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

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

# FiftyOne をインストール
# 2022/5時点でのopencv-python-headless 4.5.5.64では動作しないので4.5.4.60を再インストール
!pip install fiftyone
!pip install opencv-python-headless==4.5.4.60

 

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

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

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

 

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

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

#split="train"にすれば学習(train)データもダウンロード可能だが、容量が20GB近いので実施しない
#サンプルコードのままsplit="validation"の検証用データを表示

import fiftyone as fo
import fiftyone.zoo as foz

# List available zoo datasets
print(foz.list_zoo_datasets())

#
# Load the COCO-2017 validation split into a FiftyOne dataset
#
# This will download the dataset from the web, if necessary
#
dataset = foz.load_zoo_dataset("coco-2017", split="validation")

# Give the dataset a new name, and make it persistent so that you can
# work with it in future sessions
dataset.name = "coco-2017-validation-example"
dataset.persistent = True

# Visualize the in the App
session = fo.launch_app(dataset)

 

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

#COCOのデータセットから700のtrainデータと300のvalデータを出力する
#COCO形式で "person", "car", "chair"の3つのラベルで出力する

classes = ["person", "car", "chair"]

train_dataset = dataset[:700]
val_dataset = dataset[700:1000]
test_dataset = dataset[1000:1500]

# COCO形式でエクスポート
train_dataset.export(
    export_dir=f"/content/data/train/",
    dataset_type=fo.types.COCODetectionDataset,
    split="train",
    classes=classes,
)
val_dataset.export(
    export_dir=f"/content/data/val/",
    dataset_type=fo.types.COCODetectionDataset,
    split="val",
    classes=classes,
)
test_dataset.export(
    export_dir=f"/content/data/test/",
    dataset_type=fo.types.COCODetectionDataset,
    split="test",
    classes=classes,
)

 

YOLOX-Tinyでの学習

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

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

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

 

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

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

# YOLOX用にフォルダ作成・データセットをコピー
!mkdir dataset
!mkdir dataset/images
!mkdir dataset/images/train2017
!mkdir dataset/images/val2017
!mkdir dataset/annotations

!cp -rf /content/data/train/data/*.jpg dataset/images/train2017
!cp -rf /content/data/val/data/*.jpg dataset/images/val2017
!cp -rf /content/data/train/labels.json dataset/annotations/train_annotations.json
!cp -rf /content/data/val/labels.json  dataset/annotations/validation_annotations.json

 

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

#YOLOXをインストール
!git clone https://github.com/Megvii-BaseDetection/YOLOX.git
%cd YOLOX
!pip3 install -v -e .  # or  python3 setup.py develop

 

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

#YOLOX-tinyを使用
!cp /content/YOLOX/exps/default/yolox_tiny.py /content/YOLOX

 

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

#yolox-tiny.pyにデータセットのパスを指定する。
#/content/YOLOX/yolox_tiny.pyを開いて下記を追記する

        # Define yourself dataset path
        self.data_dir = "/content/dataset/images"
        self.train_ann = "/content/dataset/annotations/train_annotations.json"
        self.val_ann = "/content/dataset/annotations/validation_annotations.json"

        self.num_classes = 3

 

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

#yolox-tinyで学習させる(約2時間ぐらいかかる)
!python tools/train.py \
    -f yolox_tiny.py \
    -d 1 \
    -b 16 \
    --fp16 \
    -o

 

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

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

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

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

 

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

#後のコードでOPENCVのエラーが出たので一度最新をインストール
!pip uninstall opencv-python;
!pip install opencv-python

 

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

#テスト用の画像データのパスを貼り付けて、学習させたモデルでテスト(demo.py)する
TEST_IMAGE_PATH = "/content/data/test/data/000000121673.jpg"
MODEL_PATH = "/content/YOLOX/YOLOX_outputs/yolox_tiny/best_ckpt.pth"

!python tools/demo.py image \
    -f yolox_tiny.py \
    -c {MODEL_PATH} \
    --path {TEST_IMAGE_PATH} \
    --tsize 416 \
    --save_result \
    --device gpu

 

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

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

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

#テスト(demo.py)して画像認識させた結果を表示する
#下記パスはテストの日付入った一例なので適宜修正すること
from PIL import Image

OUTPUT_IMAGE_PATH = "/content/YOLOX/YOLOX_outputs/yolox_tiny/vis_res/2022_05_21_12_43_34/000000121673.jpg" 
Image.open(OUTPUT_IMAGE_PATH)

 

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

 

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

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

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

#Googleドライブの作業用フォルダの作成・マウント
from google.colab import drive 
drive.mount('/content/drive')
#!mkdir -p '/content/drive/My Drive/work/'
%cd '/content/drive/My Drive/work/'
#Googleドライブにコピー
cp -rf /content/YOLOX/YOLOX_outputs /content/drive/MyDrive/work

 

参考記事

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

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

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

 

GitHubのGist

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

ご自由にお使い下さい。

まとめ

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

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

コメント