【要約】Transformerを用いた物体検出モデル「End-to-End Object Detection with Transformers」

##はじめに
「End-to-End Object Detection with Transformers」(DETR) が気になったので、論文を読んで少し動作確認もしてみました。簡潔に記録として残しておきます。
[論文, Github]

##DETRとは（要約）
・Facebook AI Researchが今年5月に公開したモデル

・自然言語処理分野で有名なTransformerを初めて物体検出に活用

・下図のように、CNN + Transformerのシンプルなネットワーク構成

・NMSやAnchorBoxのデフォルト値等、人手による調整が必要な部分を排除し「End-to-End」な物体検出を実現

・上記を実現するためのポイントとして、"Bipartite Matching Loss"と"Parallel Decoding"の効果を主張

・物体検出だけでなく、セグメンテーションタスクへも適用可能

##推論コード例

論文よりコードを引用します。以下のように、モデル定義から推論処理までが40行程度でシンプルに書けます。

import torch
from torch import nn
from torchvision.models import resnet50

class DETR(nn.Module):

    def __init__(self, num_classes, hidden_dim, nheads,
                 num_encoder_layers, num_decoder_layers):
        super().__init__()
        # We take only convolutional layers from ResNet-50 model
        self.backbone = nn.Sequential(*list(resnet50(pretrained=True).children())[:-2])
        self.conv = nn.Conv2d(2048, hidden_dim, 1)
        self.transformer = nn.Transformer(hidden_dim, nheads,
                                          num_encoder_layers, num_decoder_layers)
        self.linear_class = nn.Linear(hidden_dim, num_classes + 1)
        self.linear_bbox = nn.Linear(hidden_dim, 4)
        self.query_pos = nn.Parameter(torch.rand(100, hidden_dim))
        self.row_embed = nn.Parameter(torch.rand(50, hidden_dim // 2))
        self.col_embed = nn.Parameter(torch.rand(50, hidden_dim // 2))

    def forward(self, inputs):
        x = self.backbone(inputs)
        h = self.conv(x)
        H, W = h.shape[-2:]
        pos = torch.cat([
            self.col_embed[:W].unsqueeze(0).repeat(H, 1, 1),
            self.row_embed[:H].unsqueeze(1).repeat(1, W, 1),
        ], dim=-1).flatten(0, 1).unsqueeze(1)
        h = self.transformer(pos + h.flatten(2).permute(2, 0, 1),
                             self.query_pos.unsqueeze(1))
        return self.linear_class(h), self.linear_bbox(h).sigmoid()

detr = DETR(num_classes=91, hidden_dim=256, nheads=8,
            num_encoder_layers=6, num_decoder_layers=6)
detr.eval()
inputs = torch.randn(1, 3, 800, 1200)
logits, bboxes = detr(inputs)

※この論文のコードのままだと、公開されている学習済みモデルをloadする事はできません。(model定義の書き方が違うため)
実際に学習済みモデルで検出を行う場合は、本家Githubのdetr_demo.ipynbのようにすると良いです。

##動作確認
学習済みモデルを使って実際に動作を確認してみました。環境は以下の通りです。
・OS : Ubuntu 18.04.4 LTS
・GPU : GeForce RTX 2060 SUPER(8GB) x1
・PyTorch 1.5.1 / torchvision 0.6.0

検出処理の実装は本家のdetr_demo.ipynbを参考にし、OpenCVでwebカメラからキャプチャした画像に対して検出を行いました。使用したモデルはResNet-50ベースのDETRです。

以下は実際の検出結果です。正常に検出が実行されていることが確認できました。(COCOのクラスに含まれない物も撮影対象にしてしまっていますが)

自分の環境では、推論処理自体は45msec（約22FPS）程度で回っているようでした。
※文献値は、V100上で28FPS

##おわり
DETRの勉強と動作確認を行いました。新しいタイプの手法なので、精度や速度の面でこれからさらに発展していくのかもと思います。
Transformerは自然言語処理専用というようなイメージを持っていましたが、最近では画像を扱うモデルへの導入も進んでいるのですね。今後、もっと勉強していきたいと思います。(Image GPTも動かしてみたいです。)

##参考

・「DETR」Transformerの物体検出デビュー
https://medium.com/lsc-psd/detr-transformer%E3%81%AE%E7%89%A9%E4%BD%93%E6%A4%9C%E5%87%BA%E3%83%87%E3%83%93%E3%83%A5%E3%83%BC-dc18e582dec1
・End-to-End Object Detection with Transformers (DETR) の解説
https://qiita.com/sasgawy/items/61fb64d848df9f6b53d1
・Transformer を物体検出に採用！話題のDETRを詳細解説！
https://deepsquare.jp/2020/07/detr/