YOLOをファインチューニングし、少量データセットで物体検知をカスタム

手持ちの少量のデータセットで、YOLO v3をファインチューニングし、カスタムした物体検知を行ないます。今回は、WHILL Model Cを画像から検知してみました。

KauzmichiShirai

YOLOv3で物体検知する方法を紹介しました。
Google Colab上でYOLO v3を使って、手持ちの画像の物体検知をしてみた

手持ちの少量のデータセットで、YOLOをファインチューニングし、カスタムした物体検知を行ないます。
今回は、WHILL Model Cを画像から検知してみました。
ファインチューニングは、学習済みのモデルの、ある畳み込みブロックの出力側の層を学習させ、入力側の層はパラメーターを凍結させます。
学習済みのモデルの表現の一部を新しい問題に適合させます。
ファインチューニングを使うことで、少量のデータセットでも性能を出せます。

以下のような概要で行ないました。

  1. iCrawlerを使って画像を集める
  2. LabelImgでアノテーションする
  3. darknet YOLOの設定ファイルをデータに合わせて修正する
  4. Google Colabで学習させる(ファインチューニング)

1.と2.は割愛します。
WHILL Model C全体が写っている画像を80枚用意しました。

GitHubからGoogle Colabに画像や設定ファイルを持っていくので、darknet YOLOを自分のGitHubにforkしておきます。
私がforkしたレポジトリは https://github.com/KazumichiShirai/darknet です。

forkができたら、2.で作った画像とアノテーションデータを/darknet/data/whillに置きます。
ディレクトリ名whillは任意の名前で大丈夫です。


まずローカル環境でdarknet/YOLOに変更を加えます。
変更後のコードはこちらに置きました。
https://github.com/KazumichiShirai/darknet

train.txtとtest.txtの用意

https://timebutt.github.io/static/how-to-train-yolov2-to-detect-custom-objects/ を参考にPythonスクリプトを書きました。
https://github.com/KazumichiShirai/darknet/blob/master/python/divide.py
使うときは、パスとopenするファイル名(whill-train.txt, whill-test.txt)を修正してください。

以下のコマンドを実行します。

python divide.py

訓練データセットはwhill-train.txtとして生成されます。
検証データセットはwhill-test.txtとして生成されます。
この2つのファイルは/darknet/cfgに置きます。

cfgファイルの作成

darknetではバッチ数やニューラルネット構造をcfgファイルで設定します。
yolov3.cfgをベースにcfgファイルを作成します。
学習に使うwhill-frozen.cfgと検出で使うwhill.cfgを作成します。

yolov3.cfgをコピーして、whill-frozen.cfgを作成します。
yolov3.cfgから次のような変更をしました。

batchとsubdivisionsを変更。

batch=16
subdivisions=4

max_batchを変更。
max_batchの回数だけ学習を繰り返すので、max_batchが大きいと学習がなかなか終わりません。
max_batchは(検出したいcleassの数) x 2000 が目安のようです*
分類したい物体はWHILL Model Cの1つなので、2000となります。
多めに学習を繰り返すために、4000としました。

max_batches = 4000

stepsはmax_batchesの80%,90%に設定します。

steps=3200,3600

分類したい数classesとfilterを変更します。
分類したい物体は1つなので、classes=1とします。
filterは(classes + 5)*mask数で計算します。
yolov3.cfgではmask=6,7,8のようにmaskは3種類になっています。
よって、filter=18に変更します。
classesとfilterはそれぞれ3箇所書かれているので、すべて変更します。

classes=1
filters=18

今回はファインチューニングで学習しました。
darknet YOLOではstopbackward=1を使って、凍結させる層を設定できます。
こちら のように、548行目に

stopbackward=1

を追加しました。

学習を終えた重みを使って検出を実行するときには、検出用のcfgファイルが必要です。
検出用のcfgファイルでは、batchとsubdivisionsを1とします。
検出用のcfgファイルをwhill.cfgとして、whill-frozen.cfgから次のような変更をしました。
Testingの直後に書かれいてるbatch, subdivisionsのコメントアウトを外して、Training直後に書かれているbatch, subdivisionsをコメントアウトします。

[net]
# Testing
batch=1
subdivisions=1
# Training
# batch=16
# subdivisions=4
# dataファイルの作成

2つのcfgファイル(whill-frozen.cfg, whill.cfg)はdarknet/cfgに置きます。

namesファイルの作成

namesファイルには分類したいクラス名を入れます。
今回はWHILL Model Cのみなので、WHILLModelCと1行だけ書きます。
ファイル名はwhill.namesにして、darknet/dataに置きました。

WHILLModelC

dataファイルの作成

分類したいクラスの数と作成したファイルへのパスをdataファイルに記載します。

classes= 1
train  = cfg/whill-train.txt
valid  = cfg/whill-test.txt
names = data/whill.names
backup = backup

今回はWHILL Model Cのみ検出したいので、classes=1となります。
backupに指定したパスには、学習途中と学習完了時の重みが保存されます。

GPU利用有効化のためのmakefile修正

YOLOのGPU利用を有効化するために、makefileを修正します。

GPU=1

ローカル環境での作業は以上です。
変更をコミットして、Google Colabでの作業に移ります。


ここからGoogle Colabでの作業になります。 Google ColabのハードウェアアクセラレータはGPUに変更しておきましょう。
参考:Google Colab上でYOLO v3を使って、手持ちの画像の物体検知をしてみた

実行したコードはこちらにコピーしました。
whill_model_c_yolo.ipynb

ビルドと学習済み重みのダウンロード

forkしたdarknet/YOLOを、GitHubからcloneしてビルドします。

%%bash
git clone https://github.com/KazumichiShirai/darknet
cd darknet
make

ImageNetで学習済みの重みをダウンロードします。

!wget https://pjreddie.com/media/files/darknet53.conv.74

ファイルを保存するために、Google Driveをmountしておきます。

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

ファインチューニング

ダウンロードした重みからファインチューニングします。
出力されるログはlogファイルにリダイレクトさせています。

%%bash
cd ./darknet/
./darknet detector \
        train \
        cfg/whill.data \
        cfg/whill-frozen.cfg \
        /content/darknet53.conv.74 > /content/log

これで学習が終わるまで待ちます。
学習が終わるまで、4~5時間程度かかりました。
学習が終わると、backupディレクトリにwhill-frozen_final.weightsという重みが保存されています。
どの時点での重みがよさそうか、また充分な性能が出ていそうか、logファイル見ながら検証する必要がありますが、ここでは割愛します。

backupディレクトリにある重みは、Google Colabのインスタンスが落ちると消えてしまうので、すぐにGoogle Driveに保存しましょう。

%cp /content/darknet/backup/whill-* /content/gdrive/My\ Drive/tmp/

学習後の重みでテストしてみる

学習には使わなかった画像を使って、WHILL Model Cが検出できるかテストしてみます。
Google Driveに、学習には使わなかったsample.JPGを置いて、確認します。

%%bash
cd ./darknet
./darknet detector test cfg/whill.data cfg/whill.cfg /content/gdrive/My\ Drive/tmp/whill-frozen_final.weights /content/gdrive/My\ Drive/tmp/sample.JPG  

結果はpredictions.jpgとして出力されます。
predictions.jpgを表示させます。

from IPython.display import Image,display_jpeg
display_jpeg(Image('darknet/predictions.jpg'))

WHILL Model Cの検知ができました!

参考にさせていただいたサイト

Yolo-v3 and Yolo-v2 for Windows and Linux
How to train YOLOv2 to detect custom objects
Google Colab上でdarknet(YOLO)を使って物体を数える【画像認識】
YOLOオリジナルデータの学習