[논문 리뷰] End-to-End Object Detection with Transformers (DETR)

  안녕하세요. 오늘 리뷰해볼 논문은 Facebook AI에서 발표했던 End-to-End Object Detection with Transformers이라는 논문이고 DETR이라고도 불립니다. 해당 논문에서 제시된 모델인 DETR은 object detection 분야에서 transformer를 backbone으로 성공적으로 적용한 모델 중에 하나이며 이후에 다른 transformer 기반 OD 모델들에게도 영향을 주고 있습니다. 

 

https://arxiv.org/pdf/2005.12872.pdf

 

https://github.com/facebookresearch/detr

GitHub - facebookresearch/detr: End-to-End Object Detection with Transformers

 

0. Abstract

  기존의 CNN 기반 모델 object detection 모델들을 보면 모델의 성능을 높이기 위해 NMS(Non-maximum supression)나 anchor generation 같이 최대한 많은 장치들을 추가해서 task에 대한 prior knowledge 최대한 많이 녹이려는 시도들이 많습니다. 해당 논문에서 제시하는 DETR은 set-based global loss와 bipartite matching 알고리즘을 사용합니다. 이 때문에 object bounding box를 엄청 많이 예측해두고 그 안에서 최종 예측을 하는 기존 방식과는 차별화를 보이고 있습니다. 

 

  DETR은 transformer encoder와 decoder를 모두 사용하는 구조이고, 정해진 갯수 K만큼의 쿼리를 decoder에 보내서 K개 만큼의 객체를 검출해내도록 합니다. 이 때 K는 모자르지 않을 정도의 고정된 충분한 숫자로 설정해둡니다. 그리고 decoder에서 나온 K개 만큼의 결과값에서 바로 object detection box의 좌표와 class를 바로 예측합니다. 이 때 K개의 prediction이 서로 겹치지 않기 때문에 K개의 prediction을 parallel하게 동시에 계산할 수 있습니다. 구조는 transformer encoder/decoder, 그 위에 box, class를 예측하는 MLP가 전부이기 때문에 굉장히 간단하고 다른 OD 모델처럼 서드파티 라이브러리를 굳이 사용하지 않아도 되어 사용하기 편하다고 합니다. 

 

  COCO에서 학습을 하였을 때 Faster-RCNN과 얼추 비슷한 성능을 보여 새로운 baseline으로서 제안을 하였습니다. Object detection 뿐만 아니라 panoptic segmentation 같은 task에서도 같은 방식으로 쉽게 학습이 가능합니다. 

 

1. Introduction

  앞서 말했던 것처럼 최근 OD 모델들은 RoI proposal을 많이 한 후에 이를 추려내는 식으로 후처리가 많이 들어갑니다. 성능도 backbone 자체가 좋은 representation을 뽑아내는지의 여부가 아니라 이 후처리 과정에 영향을 더 많이 받는다고 합니다. 하지만 해당 모델에서는 이전 같은 방식이 아니라 적당한, 고정된 수의, 겹치지 않는 prediction (set prediction)을 하여 RoI proposal에서 추려낼 필요도 없고 그래서 surrogate task를 사용할 필요도 없습니다. DETR의 구조는 아래와 같습니다. 

 

DETR 구조 1

 

  앞서 말했듯이 transformer encoder-decoder를 backbone으로 사용하였습니다. Transformer는 CNN기반 모델과 픽셀을 처리하는 방식이 다르기 때문에 이미지를 NxN개의 패치로 먼저 patchify를 합니다. 해당 방법은 ViT(Vision Transforemr) 관련 자료를 찾아보시면 좋을 것 같습니다. (여기에도 있습니다. https://dhk1349.tistory.com/6)

[논문리뷰] AN IMAGE IS WORTH 16X16 WORDS: TRANSFORMERS FOR IMAGE RECOGNITION AT SCALE

 

  Transformer encoder/decoder가 DETR 안에서 어떤 식으로 되어있는 뒷 부분에서 다루도록 하겠습니다. 일단 위의 구조를 먼저 보면 patchify된 이미지가 transformer에 들어가면 set of box predictions(set prediction)가 결과로 나오게 됩니다. 이 각각의 box prediction들을 하나의 FFN(Feed forward network)이 일괄적으로 처리합니다. 다른 기존 모델을은 여러 개의 proposal들 중에서 후처리를 통해서 최종 proposal을 추리는 과정을 거쳐야하기 때문에 proposal들을 batch에 태워서 한 번에 계산을 할 수 없습니다. 하지만 DETR에 set prediction을 FFN에 태울 때는 바로 batch에 태워서 한 번에 계산할 수 있기 때문에 parallel하게 계산이 가능하다고 저자들은 표현합니다. 

 

  하지만 해결해야할 문제가 더 있습니다. 위의 이미지처럼 새 두 마리가 있는 이미지가 있다고 가정해보겠습니다. 해당 이미지에 대한 GT는 box 2개가 됩니다. 그런데 set of predictions(K개)의 갯수가 2개 보다 많다면 K-2개의 predictions가 남게 됩니다. 이는 나머지 K-2개의 prediction를 no object class로 예측하면 됩니다.

 

  또 하나의 문제는 학습 시에 생깁니다. 위 이미지를 보면 K개의 prediction을 하지만 실제 GT는 새 두 마리가 전부입니다. 학습을 할 때 2개의 새를 잘 맞추었는지 loss를 계산해야하는데, 예를 들어 K개의 prediction들 중에서 1, 4번째 prediction box가 새라고 예측을 했을 수도 있습니다. 아니면 2,5 번 째 prediction box가 새라고 계산을 했을 수도 있습니다. 어떤 prediction box가 어떤 객체를 검출한 것인지 단순하게 정렬을 해서 확인이 불가능합니다. 이렇게 되면 loss를 계산할 수도 없고, 모델 학습도 할 수가 없습니다. 이 때문에 set of prediction과 GT들을 일일이 매칭을 해주는 작업이 필요합니다. 이 부분은 학습이 아니라 단순 매칭 알고리즘을 통해 구현을 하였습니다 (bipartite matching). 그래서 각 prediction의 class와 box coordinate예측이 GT와 최대한 매칭이 되도록 set of predictions를 매칭하여 재정렬한다고 보시면 되겠습니다. 

 

  주요한 컨셉들은 인트로에서 다 다룬 것 같습니다. 

 

2. Related Work

Set prediction, Transformers and Parallel Decoding, Object detection에 대해 다루고 있습니다. 혹시 관심있으시면 논문을 읽어보시면 좋을 것 같습니다. 

 

3. The DETR Model

DETR에서 direct set prediction이 가능하도록 하기 위해서는 크게 아래 두 가지가 필요합니다.

 

3.1. Object Detection Set Prediction Loss

  Set prediction loss는 앞에서 이야기 했던 loss를 계산할 때 생기는 문제에 대한 부분입니다. 일단은 아래와 같은 간단한 형태의 loss 식을 사용합니다. 모델을 학습할 때 set of prediction의 갯수 K(논문에서는 N이라고 씀)를 충분히 큰 숫자로 하기 때문에 일단적으로 이미지 GT 갯수가 K보다 작습니다. 남은 자리들은 no-object class로 패딩하여 채웁니다. 

 

Match loss

  이 때 matching 알고리즘을 통해 Loss가 최대한 작아지도록 prediction을 정렬합니다. 정렬을 할 때 특정 순서에 대한 cost를 계산하게 되는데 이는 class 정확도와 box coordinate의 정확도를 가지고 계산합니다. 이 같은 방식으로 매칭을 한 후에 loss를 계산해주게 됩니다. 

 

Hungarian loss

  Loss를 계산하기 위해 loss의 앞부분인 class loss는 negative log-likelihood로 계산을 하고 hungarian loss의 뒷부분인 bounding box loss는 박스의 좌표 값 자체로 계산하는 L1 loss를 사용하였고, 이렇게만 하면 박스의 크기가 큰 박스들의 loss가 더 크게 계산되기 때문에 이를 완화하기 위해 IoU Loss를 함께 사용합니다. 

 

Bounding box loss

 

그리고 no-object의 갯수가 다른 class에 비해 훨씬 많기 때문에 class imbalance 문제를 막기 위해 이 부분도 스케일링를 했다고 합니다. 

 

(추가)
Matching을 담당하는 부분의 코드를 보면 굉장히 간단한 형태로 구현이 되어있습니다. 

예를 들어 모델이 예측하고 있는 prediction box가 K1개, GT box가 K2개 있다고 가정을 해보겠습니다. 이 때, K1개의 box와 K2개의 box 간의 거리를 모두 구하여 (K1, K2) shape을 가진 거리 값을 구하게 됩니다. 이 거리 값들을 통째로 loss로 사용하여 자연스럽게 matching이 되도록 하는 구조로 구현되어있습니다.

 

 

 

3.2. DETR Architecture

아래 그림은 DETR의 구조를 더 자세하게 그린 그림입니다. 

 

DETR 구조 2

 

  Transformer encoder에 image를 넣기 위해 (3, H_0, W_0) 크기의 이미지를 CNN에 통과해서 (2048, H, W)의 크기로 맞춰줍니다. H, W는 H_0/32, W_0/32로 했다고 합니다. 이렇게 되면 token의 갯수는 HxW개가 됩니다. 

 

Transformer Encoder

  이 후에 1x1 conv를 태워서 채널 수인 2048을 조절하여 dxHW로 맞추어줍니다. 여기에 transformer가 이미지의 위치 정보를 학습할 수 있게 하기 위해 fixed positional encoding을 각 patch에 더해주어 transformer encoder에 넣어줍니다. Encoder는 기존의 transformer encoder 구조와 크게 다르지 않습니다. 

 

Transformer Decoder

  기존의 transformer와 동일하게 encoder의 결과 값들을 decoder 중간에 feed하는 식으로 구성되어 있습니다. 그러나 기존의 transformer decoder와는 다르게 DETR의 transformer decoder에서는 masked multihead attention 대신 그냥 multihead attention layer를 사용하였습니다. 기존 transformer decoder의 masked multihead attention은 결과를 autoregressive하게 계산해야기 때문에 있지만, 여기서는 이미지의 모든 feature를 참고해서 객체 검출을 해야하기 때문에 masking을 하지 않는다고 보시면 될 것 같습니다. Decoder의 입력 값을 넣는 부분에서는 K개 만큼의 set of predictions를 뽑아내야하기 때문에 K 만큼의 object queries라는 것을 넣어줍니다. DETR의 코드를 보면 query 값은 모델을 처음 초기화할 때 nn.Embedding(n_queries, num_hidden)의 크기로 초기화해주고 학습을 하면서 점점 바뀌게 됩니다. 이는 learnt positional encoding처럼 사용됩니다. 이런 방식은 transformer 기반 모델에서 positional encoding 대신 사용되기도 합니다. 

 

DETR transformer의 구조

 

Prediction Feed-forward Networks(FFNs)

FFN은 Linear+ReLU를 3번 반복한 구조입니다. Class prediction, center coordinate of box, width of box, height of box를 예측하도록 되어있습니다. 

 

 

그 뒤

 

Experiement와 ablation 부분이 궁금하시다면 원 논문을 참고해주세요.