Zenn
🗂

Detectron2のモデルをONNXへエクスポートする方法

2023/03/21に公開
PyTorch
ONNX
detectron2
tech

はじめに

Detectron2 は Facebook Research (現 Meta Research) によって開発された物体検出とセグメンテーション用のライブラリです。

ONNX は異なる機械学習フレームワーク間の相互運用を目的としたフォーマットです。ONNX フォーマットにエクスポートされたモデルは、ONNX Runtime を使用することで学習元のフレームワークを問わず実行可能になります。

カスタムデータセットで学習させた Detectron2 のモデルを ONNX へエクスポートする方法を記します。

環境構築の注意点

学習スクリプト

  • 学習には DefaultTrainer を使用します。
    • 学習後に検証したい場合は DefaultTrainer を継承したクラスを作り、build_evaluator() メソッドを定義する必要があります。
    • 学習中に検証したい場合は、さらに cfg.TEST.EVAL_PERIOD を設定する必要があります。
  • データセットは register_coco_instances() で登録します。
  • モデルや学習の情報は get_cfg() で作成した CfgNode インスタンスに設定します。
import os

from detectron2 import model_zoo
from detectron2.config import get_cfg
from detectron2.data.datasets import register_coco_instances
from detectron2.engine import DefaultTrainer
from detectron2.evaluation import COCOEvaluator

class Trainer(DefaultTrainer):
    @classmethod
    def build_evaluator(cls, cfg, dataset_name, output_folder=None):
        return COCOEvaluator(dataset_name, output_folder)

def main():
    register_coco_instances("my_dataset_train", {}, "path/to/train/annotation/file", "path/to/train/image/dir")
    register_coco_instances("my_dataset_val", {}, "path/to/val/annotation/file", "path/to/val/image/dir")

    cfg = get_cfg()
    cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
    cfg.DATASETS.TRAIN = ("my_dataset_train",)
    cfg.DATASETS.TEST = ("my_dataset_val",)
    cfg.DATALOADER.NUM_WORKERS = 2
    cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")  # model zoo の学習済みモデルの重みで初期化
    cfg.SOLVER.IMS_PER_BATCH = 2  # バッチサイズ
    cfg.SOLVER.BASE_LR = 0.00025  # 学習率
    cfg.SOLVER.MAX_ITER = 1000    # イテレーション
    cfg.SOLVER.STEPS = []         # 学習率を変化させない
    cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128   # ROIHeadsの数
    cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1 # 背景を含まないクラス数

    cfg.TEST.EVAL_PERIOD = 100     # 100 ステップごとに評価

    os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
    trainer = Trainer(cfg)
    trainer.resume_or_load(resume=False)
    trainer.train()

if __name__ == "__main__":
    main()

テストスクリプト

  • 推論には DefaultPredictor を使用します。
import cv2
import glob
import os

from detectron2 import model_zoo
from detectron2.config import get_cfg
from detectron2.data.datasets import register_coco_instances
from detectron2.data import MetadataCatalog
from detectron2.engine import DefaultPredictor
from detectron2.utils.visualizer import Visualizer

def main():
    register_coco_instances("my_dataset_train", {}, "path/to/train/annotation/file", "path/to/train/image/dir")

    cfg = get_cfg()
    cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
    cfg.DATASETS.TRAIN = ("my_dataset_train",)
    cfg.SOLVER.IMS_PER_BATCH = 1
    cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128 # 学習時と同じ
    cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1 # 学習時と同じ
    cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5 # スコアのしきい値

    # さきほど学習したモデルの重み
    cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")

    predictor = DefaultPredictor(cfg)

    # テスト結果出力ディレクトリ
    os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)

    # テスト画像のリスト
    image_paths = glob.glob(os.path.join("path/to/test/image/dir", "*.jpg"))

    for image_path in image_paths:

        # 画像ファイル読み込み
        image = cv2.imread(image_path)

        # 推論実行
        outputs = predictor(image)

        # 結果出力
        v = Visualizer(image[:, :, ::-1], MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1.2)
        out = v.draw_instance_predictions(outputs["instances"].to("cpu"))

        output_path = os.path.join(cfg.OUTPUT_DIR, os.path.basename(image_path))
        cv2.imwrite(output_path, out.get_image()[:, :, ::-1])

if __name__ == "__main__":
    main()

エクスポート方法

Docs » Tutorials » Deployment によると export_model.py スクリプトを使うと torchscript や ONNX にエクスポート可能とのことです。
- ONNX については Detectron 2 Mask R-CNN R50-FPN 3x in TensorRT が参考になります。
- 入力画像サイズは INPUT.MIN_SIZE_TEST と INPUT.MAX_SIZE_TEST および --sample-image で指定した画像から計算されるようです。
- INPUT.MIN_SIZE_TEST と INPUT.MAX_SIZE_TEST のデフォルト値は detectron2/config/defaults.py と思われます。
- TensorRT を使用する場合は入力サイズが 32 の倍数である必要がありますが、ONNX のみ使用する場合は問題ないようです。

python detectron2_repo/tools/deploy/export_model.py \
    --sample-image <path/to/sample/image> \
    --config-file detectron2_repo/configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml \
    --export-method tracing \
    --format onnx \
    --output <path/to/output/directory> \
    MODEL.WEIGHTS <path/to/model_final.pth> \
    MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE 128 \
    MODEL.ROI_HEADS.NUM_CLASSES 1 \
    MODEL.DEVICE cuda

動作確認

  • ONNX モデルの推論には onnxruntime を使用します。
  • モデルへの入力は画像データです。RGB か BGR かは未確認です。
  • モデルからの出力は順に「Box の座標・クラス・マスク・スコア・画像サイズ」となっているようです。
import glob
import os

import cv2
import numpy as np

import onnxruntime

def main():
    session = onnxruntime.InferenceSession(
        "path/to/model.onnx",
        providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])

    for input_ in session.get_inputs():
        print(input_.name, input_.shape, input_.type) 

    for output in session.get_outputs():
        print(output.name, output.shape, output.type)

    image_paths = glob.glob(os.path.join("path/to/test/image/dir", "*.jpg"))

    for image_path in image_paths:
        x = cv2.imread(image_path)
        x = cv2.resize(x, (<width>, <height>))
        x = x.astype(np.float32)
        x = x.transpose(2, 0, 1) # HWC -> CHW

        input_name = session.get_inputs()[0].name
        output_names = [output.name for output in session.get_outputs()]
        output = session.run(output_names, {input_name: x})

        print(output)

if __name__ == "__main__":
    main()

マスクの後処理

ONNX に変換したモデルから出力されるマスクの次元は (N, 1, M, M) になっています。N は検出物体の個数、M はマスク一辺の長さでデフォルトでは 28 です。入力画像と同じサイズのマスク画像を得たい場合、例えば次のような後処理が考えられます。

pred_boxes = results[0]  # shape (N, 4)
pred_masks = results[2].squeeze(axis=1)  # shape (N, 1, M, M) -> (N, M, M)
input_size = results[4]  # shape (2,)

masks = []
for pred_box, pred_mask in zip(pred_boxes, pred_masks):
    x1, y1, x2, y2 = pred_box.astype(np.int32).tolist()

    full_mask = np.zeros(input_size, dtype=np.uint8)
    resized_pred_mask = cv2.resize(pred_mask, (x2-x1, y2-y1))
    full_mask[y1:y2, x1:x2][resized_pred_mask > 0.5] = 255
    masks.append(full_mask)

参考記事

Discussion