Google ColabでMMDetectionを試す

はじめに

MMDetectionを使うことで、色々な物体検知手法を試したり、実際に学習してONNX形式で出力することが可能です。
使い方は非常に簡単で公式ドキュメント通りにやればいいのですが、Google Coalbで動かしてみたのでその手順を残します。手順自体にほとんど違いはありませんでしたが、本筋と異なるcocoデータセットのダウンロードと解凍でハマったのでそちらについても記載します。

• https://github.com/open-mmlab/mmdetection
  • GitHub
• https://mmdetection.readthedocs.io/en/latest/
  • インストール手順はこちら
• https://github.com/iwatake2222/colaboratory_study/blob/master/mmdet/mmdet_01.ipynb
  • 実際に試したノートブック

手順

1. 環境用意

1.1 ランタイムの設定

• 学習データとしてCOCOを使用します。50GByteくらい必要なのでGoogleドライブに保存します。必要に応じて容量追加してください。
• 以下のランタイム設定でGoogle Colabを起動します
  • GPU
  • Python3

1.2 MMDetectionのインストール

現在(2020年6月)、Google Colabにはデフォルトでpytorch(1.5.0)がインストールされているため必要なのはmmdetのインストールのみです。

公式手順(https://mmdetection.readthedocs.io/en/latest/install.html#a-from-scratch-setup-script )の通りにインストールするだけです。

セルの実行が完了したらランタイムを再起動します。(セルの実行画面の最後に出てくるRESTARTボタンをクリックします。)

HOME_PATH = "/content"
%cd "$HOME_PATH"
!git clone https://github.com/open-mmlab/mmdetection.git
%cd mmdetection
!pip install -r requirements/build.txt
!pip install "git+https://github.com/open-mmlab/cocoapi.git#subdirectory=pycocotools"
!pip install -v -e .  # or "python setup.py develop"

1.3 推論処理のテスト

Pretrainedモデルを使って、推論処理を試してみます。公式ガイド通りにfaster_rcnnで試します。
各モデルのチェックポイントへのリンクが https://github.com/open-mmlab/mmdetection/tree/master/configs にありますので、mmdetection/checkpoints にダウンロードします。

!mkdir -p "$HOME_PATH"/mmdetection/checkpoints
%cd "$HOME_PATH"/mmdetection/checkpoints
!wget https://open-mmlab.s3.ap-northeast-2.amazonaws.com/mmdetection/v2.0/faster_rcnn/faster_rcnn_r50_fpn_1x_coco/faster_rcnn_r50_fpn_1x_coco_20200130-047c8118.pth

demo/image_demo.py の中身を実行します。すると検知枠が付いた結果が表示されます。

%cd "$HOME_PATH"/mmdetection
from argparse import ArgumentParser
from mmdet.apis import inference_detector, init_detector, show_result_pyplot
model = init_detector("configs/faster_rcnn/faster_rcnn_r50_fpn_1x_coco.py", "checkpoints/faster_rcnn_r50_fpn_1x_coco_20200130-047c8118.pth", device="cpu")
result = inference_detector(model, "demo/demo.jpg")
show_result_pyplot(model, "demo/demo.jpg", result, score_thr=0.1)

2. COCO データセットを用意する

2.1 ダウンロード

以下のようにcocoデータセットをダウンロードします。
全部で25GByte程になります。Google Colab上にダウンロードすると後で解凍するときに容量が足りなくなるので、Google Drive上にダウンロードします。

from google.colab import drive 
drive.mount("/content/drive")
HOME_PATH = "/content"
DATASET_PATH = "/content/drive/My Drive/data/coco"
!mkdir -p "$DATASET_PATH"
%cd "$DATASET_PATH"
!wget http://images.cocodataset.org/annotations/annotations_trainval2017.zip
!wget http://images.cocodataset.org/zips/train2017.zip
!wget http://images.cocodataset.org/zips/val2017.zip
!wget http://images.cocodataset.org/zips/test2017.zip

2.2 ダウンロード完了を待つ

注意
上記コマンドによって、ダウンロード先はGoogle Driveになるのですが、どうやら一度Google Colab上に落としたものが移動されるようです。実際、ダウンロード中、利用可能なディスク容量が減っていきました。
セルの実行が終わったとしてもまだ移動が完了していない可能性があります。Google Colab上の「利用可能なディスク容量」が回復するまで待ちます。

2.3 データセットの解凍

以下のコマンドで解凍します。

注意
こちらも、ダウンロードと同様に一度Google Colab上に解凍されたものがGoogle ドライブに移動されるようです。train2017.zipとそれ以外は別々に実行した方が良いです。

%cd "$DATASET_PATH"
!unzip -n annotations_trainval2017.zip
!unzip -n train2017.zip
!unzip -n val2017.zip
!unzip -n test2017.zip

2.4 解凍を頑張る

train2017.zipの解凍中、何度もタイムアウトエラーやIOエラーが発生します。
-n オプションを付けて地道に繰り返します。

2.5 解凍結果を確認する

以下コマンドでzip内のファイル数と、実際に解凍されたファイル数を確認します。

!unzip -Z train2017.zip > train2017.txt
!ls train2017 > train2017_ls.txt
!unzip -Z val2017.zip > val2017.txt
!ls val2017 > val2017_ls.txt
!unzip -Z test2017.zip > test2017.txt
!ls test2017 > test2017_ls.txt

2.6 データセットの場所を指定

データセットの場所はconfigファイルで変更できるのですが、標準では以下のようになっている必要があります。この構造に合わせてGoogle Colab上のCOCOデータセットへのシンボリックリンクを作成します。

mmdetection
├── mmdet
├── tools
├── configs
├── data
│   ├── coco
│   │   ├── annotations
│   │   ├── train2017
│   │   ├── val2017
│   │   ├── test2017
│   ├── cityscapes
│   │   ├── annotations
│   │   ├── leftImg8bit
│   │   │   ├── train
│   │   │   ├── val
│   │   ├── gtFine
│   │   │   ├── train
│   │   │   ├── val
│   ├── VOCdevkit
│   │   ├── VOC2007
│   │   ├── VOC2012

%cd "$HOME_PATH"/mmdetection
!mkdir -p data
!ln -s "$DATASET_PATH"  data/coco

3 COCOのテスト用データで推論テスト

学習の前に、テストデータに対して推論処理をしてみます。
学習方法などを指定したconfigファイル内にデータセットの場所が指定されているので、このコンフィグファイルと先ほどダウンロードしたプリトレインモデルを指定します。

!mkdir -p "$HOME_PATH"/result_images
%cd "$HOME_PATH"/mmdetection
!python tools/test.py \
    configs/faster_rcnn/faster_rcnn_r50_fpn_1x_coco.py \
    checkpoints/faster_rcnn_r50_fpn_1x_coco_20200130-047c8118.pth \
    --show-dir "$HOME_PATH"/result_images

全画像に対して処理すると時間がかかるので、適当なところで止めます。
指定したディレクトリに結果を描画した画像が出力されます。

4. COCOの学習用データで学習

テスト用に試したのと同じconfigファイルを指定して、学習用スクリプトを実行します。

!mkdir -p "$HOME_PATH"/work_train
%cd "$HOME_PATH"/mmdetection
!python tools/train.py \
    configs/faster_rcnn/faster_rcnn_r50_fpn_1x_coco.py \
    --work-dir "$HOME_PATH"/work_train

5. onnx形式のモデルを出力する

使いやすくするために、ONNXに変換します。

!pip install onnx
%cd "$HOME_PATH"/mmdetection
!python tools/pytorch2onnx.py \
    configs/faster_rcnn/faster_rcnn_r50_fpn_1x_coco.py \
    checkpoints/faster_rcnn_r50_fpn_1x_coco_20200130-047c8118.pth \
    --out faster_rcnn_r50_fpn_1x_coco.onnx