Pytorch1.1+MaskRCNNで自前データ訓練(1)

はじめに(※若干釣りタイトル)

この記事はmask-r-cnnfaster-r-cnnで自前データ訓練をさせた際の忘備録です.※
当方環境はこんな感じです.

• OS: Windows10 pro
• Python: 3.6.6(pip管理)
• CUDA: 10.1
• Pytorch: 1.1.0
• GPU: GeForce RTX2080 x 1
• Nvidia driver: 430.39

おそらくWindowsのみで有効/必要な部分には（Windows専用）と記載していますが、別プラットフォームでの検証はしていません.

※最初に「mask-rcnnで..」とか書いていましたが, 実際にはfaster-rcnnが正しいです.mask-r-cnnを使う場合, 訓練データ側にセグメンテーション情報が必要なのですが, 今回使ったデータはBoundingBoxのみ取得可能のためです.

記載の動機

• 精度優先（YOLOとかの速度優先ではない方向）ではMaskRCNNが良いらしい
• Facebook Research MaskRCNN-benchmark(コミットID:1127bdd) 使ってみた
• とりあえずDemoで自前画像に試したところ良さげ
• ただ, Demoで示されてる使い方は「すでにある学習済み結果を使い, 解析情報を重ねた画像を作って表示する」というもの
  • ボックス座標とかcategoryとかをプログラムで受けて処理したいが, この方法が不明
  • 自前データで学習する手順が動かせるかどうか不明
• README.mdに記載された方法で実行できるのは確認できたが, 色々嵌ったので記録残しておきたい （間違いあったらツッコミほしい）

確認した手順

1. MaskRCNN実装をgit cloneする

git clone https://github.com/facebookresearch/maskrcnn-benchmark

2. 必要パッケージをインストール

requirements.txtに記載されていたのはninja, yacs, cython, matplotlib, tqdmでした.
MaskRCNN実装のINSTALL.mdにはopencv-pythonも書かれているので, 追加で必要かもしれません.
まぁ大抵他でインストール済みでしょうけど

cd maskrcnn-benchmark
python -m pip install -r requirements.txt

3. COCO Dataset API(pycocotools), NVIDIA Pytorch拡張(apex)インストール

デモ動作（maskrcnn-benchmark/demo/Mask_R-CNN_demo.ipynb)を見るだけなら要らない模様ですが,どのみち必要になるので入れてしまいましょう. maskrcnn-benchmark/INSTALL.mdに記載されているまんまです.
手順後, Jupyter Noteとかでimport apex, from pycocotools.coco import COCOとか実行して通ればOKです.

git clone https://github.com/cocodataset/cocoapi.git
cd cocoapi/PythonAPI
python setup.py build_ext install

git clone https://github.com/NVIDIA/apex.git
cd apex
python setup.py install --cuda_ext --cpp_ext

4. Mask R-CNNビルド

Windowsだと「THCCeilDiv()知らんよ」的なエラーが出てビルド失敗します.
Ad-hoc修正ですが, 以下の内容をmaskrcnn-benchmark/maskrcnn_benchmark/csrc/cuda以下の

• ROIAlign_cuda.cu
• ROIPool_cuda.cu
• SigmoidFocalLoss_cuda.cu

に導入することでビルドが通るようになります.
(Windows専用)

+ static inline long ceil_div(long a, long b)
+ {
+    return (a + b - 1) / b;
+ }
+ #define THCCeilDiv(a,b) ceil_div(a,b)

修正後にビルドコマンド実行すると, from maskrcnn_benchmark import .. が実行できるようになります.

なおPytorch側を更新した場合, 再度ビルドが必要になるのでご注意ください. 元々この記事はPytorch1.0で書こうとしてたのですが, 投稿前にPytorch1.1.0に更新したところ, 色々動かなくなって焦りました...

cd maskrcnn-benchmark
python setup.py build develop

5. demoで動作確認

動作確認にはmaskrcnn_benchmark/demo/Mask_R-CNN_demo.ipynbが動くか見るのが一番手っ取り早いです.
コードセル[7]にある image = load(...)の部分を以下のように変えてやると, 自前画像で試せます.
cv2.resize()を呼んでいる理由ですが, 画像サイズが大きすぎると(元々は4000x6000くらいのイメージ), 解析結果表示の線が相対的に細くなってしまい,
結果「DemoそのままだとPersonとかDogとか表示されてるのに, 自前画像だとcategory出ない!?」というどうでも良いところで悩んだので, 同じように悩まないよう記載しています.

import cv2
#image = load("http://farm3.staticflickr.com/2469/3915380994_2e611b1779_z.jpg") #元のコード
image = cv2.imread(自前画像へのパス)
image = cv2.resize(image,(900,600))

上が適用前, 下が適用後です. ちょっと灯篭や鳥居と人間が融合してる部分はありますが, ほぼ完璧ですね.

6. 自分で訓練

訓練用にCOCOデータセット落としておきます. 当方は"2017 Trainimages" "2017 Val images" "2017 Train/Val annotations"を使いました.
解凍してmaskrcnn-benchmark/datasets/coco/train2017/xxx.jpg, maskrcnn-benchmark/datasets/coco/annotations/instances_train2017.json 等なるよう配置しておきます.
README.mdに記載されている通り, デフォルトだと8GPUで訓練, とかになってるので

• SOLVER.IMS_PAR_BATCH バッチ当たり画像数
• SOLVER.BASE_LR 学習率

は変えてやる必要があります. (変えないとGPUメモリ足りない)
具体的には以下で実行しました. この設定ファイル, デフォルトの実行数(90000iter)だと当方環境では大体12時間かかります.実行中にはmodel_xx.pthという名前でバックアップを取ってくれるので, 次回同じパラメタで実行する場合には途中再開できるようになっています.

tools\train_net.py --config-file configs/e2e_mask_rcnn_R_50_FPN_1x.yaml DATASETS.TRAIN "(\"coco_2017_train\", \"coco_2017_val\")" SOLVER.IMS_PER_BATCH 2 SOLVER.BASE_LR 0.0025 DATASETS.TEST "(\"coco_2017_val\",)" TEST.IMS_PER_BATCH 1

注意点として, この実行時に「utf-8に属さない文字コードがうんたら」というエラーが（おそらくWindows環境依存）出ることがあり
以下の対応が必要になりました. (Pytorch側の修正)

$(USERS)/AppData/Local/Programs/Python/Python36/Lib/site-packages/torch/utils/collect_env.py

if PY3:
-        output = output.decode("utf-8")
-        err = err.decode("utf-8")
+        output = output.decode("gbk")
+        err = err.decode("gbk")

7. 訓練した結果を使って自分で推論

maskrcnn_benchmark/tools/train_net.pyでは最後にrun_test()という関数で推論を行っていますが, 大まかな流れは以下の通りでした.
なので, 例えばJuypter等からBBOX座標を使いたい場合, この流れを再現してやればよいことになります。

(以下は疑似コード)

#モデル構築
cfg = "configs/e2e_mask_rcnn_R_50_FPN_1x.yaml" #一例
model = maskrcnn_benchmark.modeling.detector.build_detection_model(cfg)
model.eval()

#推論実行
with torch.no_grad():
    images.to(device)
    output=model(images)
    output=output.to(torch.device("cpu"))

#推論結果→スコア/BBOX等変換
prediction = output.resize((image_width, image_height))
prediction = prediction.convert("xywh")
boxes = prediction.bbox.tolist()
scores = prediction.get_field("scores").tolist()
labels = prediction.get_filld("labels").tolist()

（以下は実際のコード）

from maskrcnn_benchmark.utils.env import setup_environment  # noqa F401 isort:skip
import argparse
import os
import torch
from maskrcnn_benchmark.config import cfg
from maskrcnn_benchmark.data import make_data_loader
from maskrcnn_benchmark.solver import make_lr_scheduler
from maskrcnn_benchmark.solver import make_optimizer
from maskrcnn_benchmark.engine.inference import inference
from maskrcnn_benchmark.engine.trainer import do_train
from maskrcnn_benchmark.modeling.detector import build_detection_model
from maskrcnn_benchmark.utils.checkpoint import DetectronCheckpointer
from maskrcnn_benchmark.utils.collect_env import collect_env_info
from maskrcnn_benchmark.utils.comm import synchronize, get_rank
from maskrcnn_benchmark.utils.imports import import_file
from maskrcnn_benchmark.utils.logger import setup_logger
from maskrcnn_benchmark.utils.miscellaneous import mkdir
import torchvision.transforms as T

from typing import List
import numpy as np
import cv2

def load_model(cfg:str, local_rank:List[int], distributed:bool):
    model = build_detection_model(cfg)
    device = torch.device(cfg.MODEL.DEVICE)
    model.to(device)
    optimizer = make_optimizer(cfg,model)
    scheduler = make_lr_scheduler(cfg,optimizer)
    arguments = {}
    arguments["iteration"] = 0
    output_dir = cfg.OUTPUT_DIR
    save_to_disk = get_rank() == 0
    checkpointer = DetectronCheckpointer(
        cfg, model, optimizer, scheduler, output_dir, save_to_disk
    )
    extra_checkpoint_data = checkpointer.load(cfg.MODEL.WEIGHT)
    arguments.update(extra_checkpoint_data)
    return model

def build_transform(cfg:str,min_image_size:int):
    if cfg.INPUT.TO_BGR255:
        to_bgr_transform = T.Lambda(lambda x: x * 255)
    else:
        to_bgr_transform = T.Lambda(lambda x: x[[2, 1, 0]])
    normalize_transform = T.Normalize(
        mean=cfg.INPUT.PIXEL_MEAN, std=cfg.INPUT.PIXEL_STD
    )
    transform = T.Compose(
        [
            T.ToPILImage(),
            T.Resize(min_image_size),
            T.ToTensor(),
            to_bgr_transform,
            normalize_transform,
        ]
    )
    return transform

def apply_model(cfg:str, model, original_image:np.array, min_image_size:int):
    transform = build_transform(cfg,min_image_size)
    image = transform(original_image)
    images = image.to(torch.device(cfg.MODEL.DEVICE))
    output = model(images)
    output = [o.to(torch.device("cpu")) for o in output]
    return output

num_gpus = int(os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1
cfg.merge_from_file(MASKRCNN_CFG)
cfg.merge_from_list(["DATASETS.TRAIN", "(\"coco_2017_train\", \"coco_2017_val\")",
                    "SOLVER.IMS_PER_BATCH", "2",
                    "SOLVER.BASE_LR", "0.0025",
                    "DATASETS.TEST", "(\"coco_2017_val\",)",
                    "TEST.IMS_PER_BATCH", "1"])
cfg.freeze()

model = load_model(cfg, 0, False).eval()
image = cv2.imread(PATH_TO_MY_IMAGE)
boxlist = apply_model(cfg, model, image ,1024)

ここで得られたBoxListは以下のような形で利用できます.

import pandas as pd

for field in boxlist[0].fields():
    flist = boxlist[0].get_field(field)

h,w,c = image.shape
boxitem = boxlist[0].resize((w,h))
boxitem = boxitem.convert("xywh")

scores = boxitem.get_field("scores").tolist()
labels = boxitem.get_field("labels").tolist()
boxl = boxitem.bbox.tolist()
df = pd.DataFrame({"score":scores, "label":labels, "bbox":boxl})
display(df[df.score > 0.9])

	score	label	bbox
0	0.990174	1	[1844.93017578125, 2415.5654296875, 266.969726...
1	0.996884	1	[3267.51220703125, 2431.8662109375, 244.175048...
2	0.997337	1	[2484.537841796875, 2535.302490234375, 235.046...
3	0.994330	1	[2198.98291015625, 2457.77001953125, 266.01562...
..

この情報を使って自前でバウンディングボックス描いてやるとこんな感じになりました. 上のdemo画像と大体同じ結果が得られていますね.

8. 自前データで訓練

基本的にはREADME.mdに書かれてる通りなのですが, 何点か気にする必要がありました. 改めて詳細書きます.

• MyDataset.__init__()にはtransforms引数が必須
• __getitem__()での1番目の返り値はPillow形式イメージじゃないと駄目(numpyでもtorch.tensorでもダメ）
• DataLoaderはmaskrcnn_benchmark.data.make_data_loader()を通して生成しないと色々不整合が出る
  • 今回の一番の嵌りポイント. torch.utils.data.DataLoader(MyDataset...)だと駄目です.
• DataLoaderがBounding Boxしか返さない場合, "configs/xxx_mask_rcnn_xxxx.yaml"は使用不可, "configs/xxx_faster_rcnn_xxxx.yaml"を使う
  • xxx_mask_rcnn_xxxを使うにはセグメンテーション情報を返す必要があるのですが, README.mdには見当たらないですね.
• MyDataset.__len__()も実装必要