Segmenter: Transformer for Semantic Segmentation 리뷰

Strudel, Robin, Ricardo Garcia, Ivan Laptev, and Cordelia Schmid. "Segmenter: Transformer for semantic segmentation." In Proceedings of the IEEE/CVF international conference on computer vision, pp. 7262-7272. 2021. 

ICCV 2021 Open Access Repository

 

Abstract

이미지 패치단위를 이용한 segmentation은 모호할 수 있으며 label과 일치시키기 위해서는 contextual information이 요구된다. 우리는 ViT기반의 semantic segmentation 모델 Segmenter을 제안한다. 이미지 패치로부터 임베딩 벡터를 획득하고 point-wise linear decoder 혹은 mask transformer decoder를 이용해 클래스 라벨을 획득한다. 사전훈련 모델을 사용하여 fine tuning하였을 때 ADE20K, Pascal Context에서 매우 좋은 성능을 보였다.

 

Introduction

segmentation을 하기위한 전형적인 접근 방법은 convolution을 기반으로한 인코더-디코더 형태의 모델이다. 하지만 convolutional은 이미지의 전역정보를 담는데 한계가 있고, segmentation은 종종 이미지의 전역정보를 토대로 패치들에 대한 라벨을 유추한다. 우리는 이 한계를 극복하기 위해, 트랜스포머 구조의 모델을 채택하여 단계적으로 이미지의 전역정보를 주입할 것이다.

 

Our approach : Segmenter

Segmenter은 트랜스포머 모델만을 이용한 인코더 디코더 형태의 모델로 patch embedding을 통해 픽셀단위로 class annotation한다. 인코딩된 이미지 패치는 point-wise linear mapping 혹은 mask transformer을 통해 디코딩 된다. 픽셀단위의 cross-entropy loss를 이용하고 end-to-end 방식으로 훈련된다.

Encoder

이미지는 흔히 ViT에서 입력되는 형태인 \(x=[x_1,...,x_N]\in R^{N\times P^2\times C}\) 의 형태로 임베딩 된다. 위치임베딩도 똑같은 형태로 맵핑된다. MSA, MLP도 vanilla transformer과 똑같은 형태로 진행되어 feature vector을 출력한다.

Decoder

인코딩된 벡터 \(z_L\in R^{N\times D}\)은 K개의 클래스를 가지는 segmentation map \(s\in R^{H\times W\times K}\)으로 디코딩된다. 디코더는 패치단위로 인코딩된 벡터를 패치단위의 클래스 점수로 학습된다. 이후 bilinear interpolation을 통해 업샘플링 된후 픽셀 단위의 점수를 반환한다.

 

Linear : 패치 인코딩 \(z_L\in R^{N\times D}\)에 대해 point-wise linear layer를 적용하여 패치단위의 class logits \(z_{lin} \in R^{N\times K}\) 으로 변환한 후 2D의 형태로 변형되어 원본 크기로 업샘플링 된다.

 

Mask Transformer : 트랜스포머 기반 디코더에 대해 K개의 클래스에 대해 \(cls = [cls_1,...,cls_K]\in R^{K\times D}\) 인 학습가능한 토큰을 생성한다. 이 토큰은 Fig2와 같이 patch encodings \(z_L\)과 함께 입력되어 K개의 마스크를 생성한다. 마스크는 2D의 형태로 변형된 후 업샘플링되고 softmax를 통해 최종 segmentation map을 위한 스코어를 반환한다.

 

모델이나 기술쪽에서 특별하게 쓰인것은 없는 것 같다.

 

자세한 실험 결과 및 ablation study는 원 논문을 참조하길 바란다.