오늘은 DeTR(Detection Transformer) 논문에서 소개한 모델을 공부해보도록 하자.
https://arxiv.org/abs/2005.12872
End-to-End Object Detection with Transformers
We present a new method that views object detection as a direct set prediction problem. Our approach streamlines the detection pipeline, effectively removing the need for many hand-designed components like a non-maximum suppression procedure or anchor gene
arxiv.org
Method
전체적인 알고리즘 개요도는 밑의 그림과 같다.
먼저 입력 이미지를 CNN에 통과시켜 C × H × W 크기의 feature map을 만든다. 일반적으로 C = 2048 정도로 설정하며, H와 W는 원본 이미지 해상도의 약 1/8 수준이다. 이후 1×1 convolution을 적용해 채널 수를 C에서 Transformer가 처리하는 hidden dimension인 d로 줄여 d × H × W 형태의 feature map으로 만든다. 이 feature map을 공간 차원에 대해 펼쳐서 총 H × W 개의 d-차원 토큰으로 변환하고, positional embedding을 더한 뒤 Transformer encoder에 입력한다.
Encoder는 이 토큰들 사이의 global self-attention을 수행하여 이미지 전체의 정보가 서로 전달되도록 만든다. 이어서 encoded token들을 Transformer decoder로 넘기면, 여기에는 학습 가능한 object query들이 입력된다. 각 query는 “하나의 object를 나타내는 latent slot”처럼 동작한다고 이해할 수 있다. Decoder에서 query들은 먼저 서로 self-attention을 수행하며 서로 다른 객체를 담당하도록 조정되고, 이후 encoder output token들에 대해 cross-attention을 수행하면서 자신이 예측해야 하는 object와 관련된 이미지 영역의 정보를 받아오게 된다. 이 과정을 통해 각 query는 하나의 object에 대응되는 표현을 얻게 되고, decoder가 끝난 뒤에는 prediction head를 통해 각 query로부터 class distribution과 bounding box 좌표를 출력한다.
일반적으로 query의 개수 N은 이미지에 존재할 수 있는 실제 object 수보다 충분히 크게 설정되기 때문에, 예측된 object와 ground-truth object 사이의 1:1 대응을 찾기 위해 Hungarian algorithm을 사용한다. 매칭되지 않은 query는 no-object 클래스를 예측하도록 supervision을 주고, 매칭된 query는 class loss와 bounding box loss 모두를 통해 학습시킨다. 이러한 방식으로 DETR은 별도의 NMS(Non Maximum Suppression) 없이 end-to-end object detection을 수행할 수 있게 된다.