mmdetのmask2formerを独自データセットで学習する

この記事は？

DL系のライブラリを使うときには以下の設定を自分の学習したいデータに合わせて調整する必要があります。mmdetではこれらを全て1つのconfigファイルから設定することができます。本記事では具体的にconfigのどこを調整すればいいのかをmask2formerを例に解説します。

• 独自データセットの指定
• 事前学習重みの指定
• 学習パラメータの設定

本記事の実行環境は以下のようになっています。

• mmdet==3.1.0
• mmengine==0.8.4
• mmcv==2.0.1
• torch==2.0.1
• torchvision==0.15.2

環境構築に使ったDockerfile

FROM pytorch/pytorch:2.0.1-cuda11.7-cudnn8-devel

ENV DEBIAN_FRONTEND noninteractive
ENV TZ "Asia/Tokyo"

RUN apt update -y && \
    apt install -y \
    curl \
    git \
    vim \
    libgl1-mesa-dev \
    libglib2.0-0

RUN pip install -U openmim
RUN mim install mmengine && \
    mim install "mmcv>=2.0.0" && \
    mim install mmdet

mmdetのmask2formerを独自データセットで学習する

独自データセットの指定

用意するデータセットのフォーマット

mmdetで独自データセットを利用するときに推奨されているのはCOCOと同じ形式でデータセットを用意することです。つまり、

• データセットの全てのアノテーションを記述したjsonファイル1つ
• データセットの全ての画像を収めたディレクトリ1つ

があれば十分です。trainデータセット以外にvalidationやtestデータセットを用意する場合は、jsonファイルをスプリットごとに用意すれば良いです。これらのファイル名/ディレクトリ名は任意なので好きな名称をつけて大丈夫です。

アノテーションのフォーマットの詳細はmmdetの公式にガイダンスがあるのでこちらを参考にすると良いかと思います。

この結果、以下のような構成で独自データセットを作成したとしましょう。

.
└── my_dataset/
    ├── annotations/
    │   ├── train.json
    │   └── val.json
    └── images/
        ├── train/
        │   ├── img1.jpg
        │   ├── img2.jpg
        │   └── ...
        └── val/
            ├── img1.jpg
            ├── img2.jpg
            └── ...

独自データセットのパスをconfigの中で指定する

今回はResNet50をバックボーンにしたinstance segmentation用のmask2formerのconfigをベースに学習用configを作っていきます。独自データのパスを指定するためのパラメータは *_dataloader と *_evaluator の変数定義部分に出てくるのでここを変えます。

変更前.py

train_dataloader = dict(
    dataset=dict(
        type=dataset_type,
        ann_file='annotations/instances_train2017.json',
        data_prefix=dict(img='train2017/'),
        pipeline=train_pipeline))
val_dataloader = dict(
    dataset=dict(
        type=dataset_type,
        ann_file='annotations/instances_val2017.json',
        data_prefix=dict(img='val2017/'),
        pipeline=test_pipeline))
test_dataloader = val_dataloader

val_evaluator = dict(
    _delete_=True,
    type='CocoMetric',
    ann_file=data_root + 'annotations/instances_val2017.json',
    metric=['bbox', 'segm'],
    format_only=False,
    backend_args={{_base_.backend_args}})

変更後.py

data_root = 'my_dataset'  # データセットのパスを指定します。データセット自体はどこに置いても良いです
train_dataloader = dict(
    dataset=dict(
        type=dataset_type,
        data_root=data_root,
        ann_file='annotations/train.json',  # data_rootからみたアノテーションファイルの相対パスを指定します
        data_prefix=dict(img='images/train/'),  # data_rootからみた画像データの相対パスを指定します
        pipeline=train_pipeline))
val_dataloader = dict(
    dataset=dict(
        type=dataset_type,
        data_root=data_root,
        ann_file='annotations/val.json',
        data_prefix=dict(img='images/val/'),
        pipeline=test_pipeline))
test_dataloader = val_dataloader

val_evaluator = dict(
    _delete_=True,
    type='CocoMetric',
    # データローダーと評価クラスが独立しているので改めてアノテーションのパスを指定する必要があります
    # CocoMetricはdata_root引数を持たないのでフルパスをで指定する必要があります
    ann_file=data_root + 'annotations/val.json',  
    metric=['bbox', 'segm'],
    format_only=False,
    backend_args={{_base_.backend_args}})

これでデータセットのパス指定は完了です。

mmdetのconfigについてのtips

mmdetは設定すべきパラメータが膨大なので、基本的に公式のサンプルを微修正して使う方法が一般的かと思います。

mmdetのconfigは継承をサポートしているため、ベースになるconfigも考慮しないと思った通りの挙動にならないことがあります。個人的なベストプラクティスとしては、最初に実験をするときには継承しているパラメータを全部1枚のconfigにまとめるようにしています。継承して綺麗にするのはある程度実験のバリエーションが出てきて、設定としてしてまとめた方がいい部分が明確になってからでも遅くないと思います。継承の旨みが出てくるのは比較実験でパラメータを一部変えて性能比較をするときなどに出てくるのかな、という印象です。

パス以外の設定

データセットのパス以外にも学習するオブジェクトのクラスに関するパラメータを独自データのものに合わせる必要があります。
元々のconfigでこれらの値がハードコードされている場合、適宜見つけて修正する必要があります。ある程度モデルだったり前処理自体の習熟度が高くないと変更箇所が自明でないので、その辺りがmmdetが敬遠される理由なのかなと思います。

今回使ったconfigの場合は以下のパラメータを変更する必要があります。

• num_thing_classes
  • オブジェクトのクラス数
• num_stuff_classes
  • panopticで使うstuffのクラス数。0にすればinstance segmentationになる
• CrossEntropyLoss のクラスごとの重み
  • 元の実装では80クラス分のリストが作られているので、これを独自データのクラス数の長さに合わせて調整する

ハードコードされている変更すべき箇所の探し方

以下の箇所は怪しいことが多いので最初にチェックします。これで漏れがある場合は最終的には目grepで気合で探します。

• lossの周り
  • lossはクラス単位で処理することが多いので出現しやすい
• データ前処理周り
  • クラスに関連する処理で出現する
• モデルのヘッド周辺
  • クラス数に応じて出力数を変えたりするので出現しやすい

COCOで学習しているモデルの場合、80 で検索をかけてハードコード部分を探すとかもよくやります。

最後にデータセットのメタ情報（クラスラベル）を渡します。これは dataset のパラメータで、これを渡さないと学習が止まります。train/val/test全ての dataset に渡す必要があります。

metainfo.py

metainfo = dict(
    classes=(
        'class1',
        'class2',
        ...,
    ),
    # paletteはクラスを描画するときの色を指定します。たぶん指定しなくても動きます
    palette=[
        (200, 200, 200),  # RGBの値を指定する
        (200, 100, 200),
        ...,
    ],
)
train_dataloader = dict(
    dataset=dict(
        type=dataset_type,
        data_root=data_root,
        metainfo=metainfo,
        ann_file='annotations/train.json',  # data_rootからみたアノテーションファイルの相対パスを指定します
        data_prefix=dict(img='images/train/'),  # data_rootからみた画像データの相対パスを指定します
        pipeline=train_pipeline))

# val/testは略

事前学習重みの指定

load_from のパラメータに事前学習重みのパスやURLを指定することでfinetuneになります。クラス数を変えてモデルアーキテクチャが変わった場合、パラメータの形状が変わった部分以外を読み込み、パラメータ形状が変わった部分に関してはランダムに初期化します。torchの重み読み込みにおいて strict=False になったときの挙動だと思えば良いです。

load_from.py

# ダウンロードした重みをcheckpoints以下に置いたので、そのパスを指定します
load_from = 'checkpoints/mask2former_r50_8xb2-lsj-50e_coco_20220506_191028-41b088b6.pth'

学習パラメータの設定

これは実験条件に合わせて適当に変えるところかなとは思います。よく変えるのは、VRAM消費量に関連するところと学習のスケジューリング部分でしょうか。以下は一例です。

• 学習中の画像サイズ
  • mmdetの学習はデフォルトだと大きいサイズを使っていることが多いので、VRAMが貧弱だとデフォルト設定だと回らないことがあります
  • そこで画像サイズを適当に落とすことが多いです。関連箇所は例によってハードコードされている場合が多いので気合いで画像サイズに関連してそうな箇所を探して直します。pipelineが前処理のキーワードなのでこの近傍に存在していることが多いです
• 学習率/epoch数
  • mask2formerの場合、optim_wrapper/param_scheduler で設定できました

ここまで設定したらあとは python tools/train.py path-to-confing を動かして祈るだけです。

最後に

mmdetは最初は相当とっつきづらいですが、慣れると「まあ拡張性高くやろうとしたらそう実装するしかないよね」みたいな実装になっていて納得はできます。ある程度知識があって拡張性高くガチャガチャやりたい人にとってはいいライブラリだと思います。ぜひ使ってみてください。