物体検知評価のためにCOCOデータセットを使う (学習編)

前回はCOCOデータセットの紹介をしました。今回はCOCOデータセットを使って実際に学習を行ってみます。

前回ご説明したようにCOCOを使った学習は、train2014の82,783枚と、val2014の40,504枚からvalidation用の5,000枚を引いた残り(val35k)をあわせた、trainval35kがよく用いられます。このデータのアノテーションデータは、instances_trainval35k.jsonと呼ばれるものですが、なぜか検索しても見つからなく、仕方ないのでまずこのアノテーションデータを作るところから始めます。

instances_train2014.jsonとinstances_valminusminival2014.jsonは入手可能ですので、instances_trainval35k.jsonはこれらをマージすればよいことになります。それぞれのjsonファイルの中身を見てみると、'info', 'licenses', 'categories', 'images', 'annotations'が共通するkeyで、このうち'info', 'licenses', 'categories'はvalueも同じでよいようです。あとは'images'と'annotations'をマージすれば良さそうなので、以下のようなスクリプトで、instances_trainval35k.jsonを作ってみました。

import json

def main():
    f0 = open('instances_train2014.json')
    d0 = json.load(f0)
    f1 = open('instances_valminusminival2014.json')
    d1 = json.load(f1)

    i0 = d0["images"]
    i1 = d1["images"]
    i0.extend(i1)
    a0 = d0["annotations"]
    a1 = d1["annotations"]
    a0.extend(a1)

    d2 = {'info': d0['info'],
          'licenses': d0['licenses'],
          'categories': d0['categories'],
          'images': i0,
          'annotations': a0}

    f2 = open('instances_trainval35k.json', 'w')
    json.dump(d2, f2)

if __name__ == '__main__':
    main()

以下のように実行すれば、アノテーションファイルが作成されます。

$ python create_trainval35k.py

次にディレクトリを作ってデータを配置します。traival35kの中にtrainとvalの全イメージを入れてしまえば、instances_trainval35k.jsonに存在するイメージだけが使われます。

  coco
      images
          trainval35k
            - COCO_train2014_000000000009.jpg
            - COCO_train2014_000000000025.jpg
            - ...
      annotations
        - instances_trainval35k.json

COCOデータの用意ができたので、まずはSSDを改造するベースとして使ってきたpytorchのSSDで、COCOデータを用いた学習をしてみます。オリジナルの実装にCOCO用の学習コードが含まれているので、それを使うだけ。 data/coco.py が、前回説明したCOCO APIを使って学習データを取り出すコードです。ご興味がある方は中身をご確認ください。

学習関連のパラメータは data/config.py にまとめられています。以下の部分が使われますが、num_classesだけ変更しています。COCOのクラス数は80 (backgorundを含めて81) だと思うので、なぜ201となっているのか不明なのですが、ひとまず81に変更します。

coco = {
    # 'num_classes': 201,
    'num_classes': 81,
    'lr_steps': (280000, 360000, 400000),
    'max_iter': 400000,
    'feature_maps': [38, 19, 10, 5, 3, 1],
    'min_dim': 300,
    'steps': [8, 16, 32, 64, 100, 300],
    'min_sizes': [21, 45, 99, 153, 207, 261],
    'max_sizes': [45, 99, 153, 207, 261, 315],
    'aspect_ratios': [[2], [2, 3], [2, 3], [2, 3], [2], [2]],
    'variance': [0.1, 0.2],
    'clip': True,
    'name': 'COCO',
}

学習は以下のようにして開始します。

$ CUDA_VISIBLE_DEVICES=0 python train.py --dataset=COCO --dataset_root=/mnt/ssd/data/coco/

学習結果の評価は次回とすることにして、その前に以前ご紹介したResNeXtをベースネットとして用いるFSSDを、COCOデータで学習してみます。train_fssd_resnext.py内でcfg = vocd512fとなっている部分を、cfg = cocoに変更します。また、data/config.py内で以下のようにします。

coco = {
    'num_classes': 81,
    'lr_steps': (100, 130, 150),
    'max_epoch': 150,
    'feature_maps': [64, 32, 16, 8, 4, 2, 1],
    'min_dim': 512,
    'steps': [8, 16, 32, 64, 128, 256, 512],
    'min_sizes': [35, 76, 153, 230, 307, 384, 460],
    'max_sizes': [76, 153, 230, 307, 384, 460, 537],
    'box': [6, 6, 6, 6, 6, 4, 4],
    'aspect_ratios': [[2,3], [2,3], [2,3], [2,3], [2,3], [2], [2]],
    'variance': [0.1, 0.2],
    'clip': True,
    'name': 'COCO',
}

以下のコマンドで学習をスタートします。

$ CUDA_VISIBLE_DEVICES=0 python train_fssd_resnext.py --dataset=COCO --dataset_root=/mnt/ssd/coco/ --loss_type=cross_entropy --weight_prefix=VOC512_FSSD_RESNEXT_ --batch_size=8 --use_pred_module=True

どちらともロスの値が下がっていくようなら、ひとまずは成功です。

NEWS & BLOGAI Solution Company, Sigfoss