mmdetectionで自作のデータセットを学習させる

手順の概要

1. データセットの用意
2. モデルを選定
3. Configファイルを書く
   1.1. 自作Configのディレクトリを作る
   1.2. Configを書いていく
   1.3. モデルのハイパーパラメータを変える
4. 学習の実行

データセットの用意

mmdetectionではCOCOフォーマットやPascalフォーマットを推奨している．
そのため，mmdetectionで学習する場合はCOCOフォーマットやPascalフォーマットのデータを用意すると良い．
自分で画像をラベリングしてデータセットを作る場合はCVATを使うと良いだろう．

今回のデータセットはインスタンスセグメンテーションを行うCOCOフォーマットのデータとする．
データセットはmmdetectionのdataディレクトリに以下のように格納されているとする．
データセットのクラス数は40とする．

mmdetection/data/mydataset
                     |---image
                     |     |---train
                     |     |     |---aaa.jpg
                     |     |     |---bbb.jpg
                     |     |     |---ccc.jpg
                     |     |---val
                     |          |---zzz.jpg
                     |          |---yyy.jpg
                     |
                     |---annotation
                             |---train
                             |　 　|---train.json
                             |---val
                                  |---val.json

モデルの選定

mmdetectionのモデルはgithubのページに一覧が書かれている．
物体検出モデルのFaster R-CNNやインスタンスセグメンテーションモデルのMask R-CNNなどがある．
この中から自分の使ってみたいモデルを選ぶことになる．
もし，何を使えばよいのかわからなかったら，Faster R-CNNやMask R-CNNを動かしてみると良い．
今回はこのMask R-CNNを使ってみることとする．

Configの書き換え

自作Configのディレクトリを作る

mmdetectionディレクトリ内で以下のコマンドを使って，新しいConfig用のディレクトリを作る．
そして，Mask R-CNNのディレクトリも作り，自分のconfigファイルも作る．

$ mkdir myconfigs
$ cd myconfigs
$ mkdir mask_rcnn
$ cd mask_rcnn
# touch mask_rcnn_r50_fpn_1x_coco.py

Configを書いていく

上で作成したmask_rcnn_r50_fpn_1x_coco.pyに以下のコードを書いていけばよい．
適宜，説明を加えた．

# mmdetectionではconfigをオーバーライドできる
_base_ = ['../../configs/mask_rcnn/mask_rcnn_r50_fpn_1x_coco.py']

# 自作のデータセットのクラスをここに書いていく
dataset_type = 'CocoDataset'
classes = ('A','B','C',...)

# データセットのクラス数を40としたため以下のように変える
# COCOのデータセットのクラス数は80だから書き換えないとエラーを吐く
model = dict(
    roi_head=dict(
        bbox_head=dict(num_classes=40,),
        mask_head=dict(num_classes=40,)))

# データセットの画像ディレクトリやアノテーションファイルの場所を書く
# 自分のデータセットディレクトリに合わせること！！
base_directory = 'data/mydataset/'
train_image = 'image/train'
train_ann = 'annotation/train/train.json'
val_image = 'image/val'
val_ann = 'annotation/val/test.json'

dataset_train = dict(
    type=dataset_type,
    classes=classes,
    img_prefix=base_directory + train_image,
    ann_file=base_directory + train_ann,
    pipeline=train_pipeline
)

dataset_val = dict(
    type=dataset_type,
    classes=classes,
    img_prefix=base_directory + val_image,
    ann_file=base_directory + val_ann,
    pipeline=test_pipeline
)

data = dict(
    samples_per_gpu=4,
    workers_per_gpu=2,
    train = dataset_train,

    val= dataset_val,
    test= dataset_val
)

モデルのハイパーパラメータを変える

上で書いたconfigの下に追記すればよい

# optimizer
# pytorchのoptimizerが使える
# typeにクラス名，後ろに引数を書く
# オーバーライドすると，不要な設定がされることがあるので，__delete__でオーバーライドをしないようにできる
optimizer = dict(type='SGD', lr=0.01, momentum=0.9, weight_decay=0.0001)
optimizer = dict(__delete__=True, type='Adam', lr=0.001)

# scheduler
# pytorchのscheculerが使えると思ってよい
# 本家の実装はwarmupなどないが，mmdetectionでは拡張されている
lr_config = dict(
    policy='step',
    warmup='linear',
    warmup_iters=500,
    warmup_ratio=0.001,
    step=[9, 11])[

# epoch数
# ファインチューニングならそこまでepoch数はいらないと思われる
runner = dict(type='EpochBasedRunner', max_epochs=15)

# モデルのパラメータを何エポックごとに保存するのか設定
# とりあえず3にしておく
checkpoint_config = dict(interval=3)

# evaluationで一番良かったモデルを保存する
evaluation = dict(metric=['bbox', 'segm'], save_best = 'auto')

# 学習データで学習した後に，検証データ検証するようにする
workflow = [('train', 1), ('val', 1)]

学習の実行

以下のコマンドで学習できる

$ python tools/train.py myconfigs/mask_rcnn/mask_rcnn_r50_fpn_1x_coco.py