YOLACTをcustom datasetで学習させる際の障壁

こんにちは。

皆さん、YOLACTをご存じでしょうか?おそらく多くの方が「YOLO」を知っているかと思います。簡単に言うとYOLACTは、YOLOのようにリアルタイムに1-stageでインスタンスセグメンテーションを行なえるものです。
こちらを研究の一部で使用するため実際に学習を試していたところ、多くの問題が発生したため、本記事ではその問題と個人的な解決を列挙してみました。

https://twitter.com/bbt_dev/status/1283906127105679360?s=20

ちなみにYOLACTの学習については、基本的に下記の記事に従ってデータセットを作成し、 qiita.com

下記の記事に従って学習を進めればよいです。 qiita.com

2つ目のQiita記事のtrain.pyの実行時にエラーが発生するため、その対処法になります。各自の環境によっては全く問題が出ない可能性がありますのでご了承ください。

ちなみに、私はラベリングの段階を「labelme」という主流のアプリの替わりに、OpenCVで自作した一物体専用領域抽出スクリプトを使用しました。これについては現在の研究内容に近く、コードの公開等がどこまで許されるかわからないため内緒です。

実行環境

Google Colab using GPU (CUDA 9.2 ?)
custom dataset (only a category)
torch==1.4.0, torchvision==0.5.0
resnet50-19c8e357.pth

train.py実行のためのコマンド

python train.py --config=yolact_base_config --save_folder=./save_folder/ --save_interval=1000 --batch_size 2

訓練画像の生成は、下記のツールの改良前のものを使用していました。 ume-boshi.hatenablog.jp



DCNv2がビルドできない

バージョンの問題のせいだと考えられます。 Google Colabが下記のようなバージョンで実行できているため、それにそろえるのが良いのではないでしょうか。

CUDA : 9.2
torch : 1.4.0
torchvision : 0.5.0

ただし、ローカルのwindowsで環境構築していた場合、python3.8のバージョンでは「torch」と「torchvision」がpipで直接インストールできなかったので、wheelファイルを指定する必要があります。

このサイトからwheelデータをDLし、コンソールで ↓ のようなコマンドを入力するとinstallできるはずです。

python -m pip install ダウンロードしてきたファイル名.whl

KeyError: Caught KeyError in DataLoader worker process 0.

./data/config.py内の"label_map"の値をうまく設定できていませんでした。

("category_id":1),  

というように記述していたのですが、

(1:1),

というように記述するのが正しいようです。 いわゆる自作のdatasetに関する設定ミスでした。

同様に、「KeyError: 0」といったエラーは、データセットを正しく読めていない可能性が高いです。

RuntimeError: CUDA out of memory.

専用GPUメモリの容量が足りていないようです。 バッチサイズを小さくすることで実行できるようになります。

製作者曰く、バッチサイズ1つにつき 1.5GBほどメモリを取るそうです。

私のノートPCに搭載されているMX230は2GBのメモリを搭載しているのですが、バッチサイズを2にしても同様のエラーが出ました。

バッチサイズを1にして実行すると、次のエラーが出るようになりました。

RuntimeError: strides[cur - 1] == sizes[cur] * strides[cur] INTERNAL ASSERT FAILED at /pytorch/torch/csrc/jit/fuser/executor.cpp:175

一番悩まされた問題です。issueを除いていても解決策らしきものが載っておらず、手探りで色々と調節していると偶然解決できました。

私の環境の答えとしては、学習データに関するjsonファイルにおける、"category_id"に"0"を指定してその周りの値が原因でした。

-Before-

train/test.json
  "categories": [{"id": 0, ...}],
  "annotations":[{"category_id": 0, ...}, ... ]

config.py
  dataset_base = Config({...
     'label_map': {0:0},
  })

-After-

train/test.json
  "categories": [{"id": 1, ...}],
  "annotations":[{"category_id": 1, ...}, ... ]

config.py
  dataset_base = Config({...
      'label_map': {1:1},
  })

↓ 初のissueへのコメント!
Internal Assertion error · Issue #275 · dbolya/yolact · GitHub
回答が間違ってたら怖いなぁ


終わりに

機械学習には疎いため1週間以上はエラーに悩まされていたのですが、学習に成功してひと安心です。

他にも沢山エラーが出た気がしますが、忘れてしまいました。なので皆さんの問題もコメントで質問していただければ、何か手伝えることがあるかもしれないのでお気軽にどうぞ。



ちなみに私はテストとして、{40枚程度, 背景シンプルな学習画像, 1時間半程度の学習時間}という条件でDualShock3を学習させてみました。その結果 ↓

f:id:ume-boshi:20200911045104p:plain
認識結果。学習画像の枚数を増やしたらさらに化けそう。


<追記> YOLACTなどのinstance segmentation用のデータセットを作成するために,1つの動画から1物体のアノテーションを楽にできるツールを作っていました。

ume-boshi.hatenablog.jp