Segmenter: Transformer for Semantic Segmentation (ICCV2021)

안녕하세요. 지난 포스팅의 [IS2D] Rethinking Semantic Segmentation from a Sequence-to-Sequence Perspective with Transformers (CVPR2021)에서는 Transformer 기반 segmentation 모델의 원조격인 SETR에 대해서 소개하였습니다. 오늘은 그 이후에 등장한 Segmenter라는 모델에 대해서 소개하고자 합니다. 

 

 

Background

기본적으로 의미론적 영상 분할 (Semantic Segmentation)은 영상 내의 각 픽셀에 대해 어떤 카테고리 레이블에 속하는 지 예측하는 dense prediction 문제라고 볼 수 있습니다. 이를 통해 사용자는 전체 영상 내에 완벽한 이해 (complete understanding)를 제공해주죠. 이러한 문제는 자율주행, 이상탐지, 그리고 의료영상 등과 같은 정말 다양한 도메인에서 필수 기술로서 자리매김하고 있습니다. 

 

이를 위해 제안된 다양한 모델들이 있습니다. 대표적으로 CNN-based 모델들은 inductive bias를 기반으로 설계되어 상대적으로 적은 양의 데이터로도 충분히 좋은 성능을 얻을 수 있다는 장점은 있으나 global information을 활용하지 못합니다. 대표적으로 FCN (CVPR2016)과 DeepLab Series (ICLR2015, TPAMI2017, arxiv2017, ECCV2018) 등이 있습니다. 물론, DeepLab Series와 같은 경우에는 ASPP와 같은 set of large kerel size convolution을 이용하여 이를 해결하고자 하였으나 여전히 convolution 연산 기반으로 동작하는 한계점이 존재합니다. 

 

본 논문에서는 이러한 한계점을 극복하기 위해 semantic segmentation을 Transformer 기반의 sequence-to-sequence problem로 해석을 하고자 하였습니다. 이를 통해 모델의 모든 부분에서 contextual information을 효과적으로 추출할 수 있게 되었죠. 하지만 입력 영상 전체를 sequence로 해석하면 픽셀 개수만큼의 token이 나와 quadratic complexity가 요구되는 Transformer 특성 상 굉장히 비효율적으로 동작하게 됩니다. 이를 위해 기존의 ViT와 마찬가지로 입력 영상을 패치단위로 처리한 뒤 Transformer Encoder로 입력합니다. 사실 이 논문에서 중요한 부분은 바로 Decoder 파트입니다. 본 논문에서는 Linear 또는 Transformer Decoder인 Mask Transformer를 제시하였으며 Mask Transformer 의 경우 Class Token을 도입하여 출력 sequence과 class token 간의 상관성을 통한 decoding 방법론을 제시하였습니다. 아래 그림은 본 논문을 통해 학습한 Segmenter의 attention map을 시각화한 결과입니다. 

 

 

 

Segmenter

1) Overall Architecture

 

 

그림 2는 Segmenter의 전체적인 개요를 보여주고 있습니다. 크게 보면 단순하게 Transformer Encoder + Mask Transformer (Decoder)로 구성되어 있기 때문에 아주 쉽게 이해하실 수 있을 것 입니다. 

 

2) Encoder

기본적으로 본 논문에서 사용하는 Encoder는 ViT에서 사용하는 Transformer Block과 동일한 구조를 사용합니다. 따라서, 이 부분은 굉장히 단순하게만 설명드리자면 Patchify + Patch Embedding (Flatten & Projection) + Positional Embedding + Transformer Encoder (MHSA & FFN)으로 구성됩니다. 여기서 Transformer Encoder 파트는 총 $L$번 반복됩니다.

 

3) Decoder

Encoder를 통해 얻은 sequence를 $\mathbf{z}_{L} \in \mathbb{R}^{N \times D}$라고 했을 때 Decoder의 역할은 segmentation map $\mathbf{s} \in \mathbb{R}^{H \times W \times K}$로 디코딩하는 것이 최종목표입니다. 여기서 $K$는 클래스의 개수를 의미합니다. 본 논문에서는 이를 위해 두 가지 scheme을 제시합니다. 각각 Point-wise Linear Layer와 Mask Transformer입니다. 

 

3-1) Point-wise Linear Layer

STEP1. Point-wise Linear Layer를 patch encoding $\mathbf{z}_{L}$에 적용하여 patch-level class logits $\mathbf{z}_{\text{lin}} \in \mathbb{R}^{N \times K}$로 매핑한다. 

 

STEP2. 2D feature map으로 reshape하여 $\mathbf{s}_{\text{lin}} \in \mathbb{R}^{H / P \times W / P \times K}$를 얻는다. 

 

STEP3. $\mathbf{s}_{\text{lin}}$에 bilinear upsampling을 통해 원본영상 크기 $\mathbf{s} \in \mathbb{R}^{H \times W \times K}$를 얻는다. 

 

STEP4. Softmax를 class dimension에 적용하여 final segmentation map을 얻는다. 

 

3-2) Mask Transformer

STEP0. $K$개의 학습 가능한 class embedding $\textbf{cls} = [\text{cls}_{1}, \text{cls}_{2}, \times, \text{cls}_{K}] \in \mathbf{R}^{K \times D}$를 초기화하고 각 토큰별로 semantic class를 배정한다. 

 

STEP1. class mask의 set을 $\text{Mask} (\mathbf{z}_{\mathbf{M}}^{'},\mathbf{c}) = \mathbf{z}_{\mathbf{M}}^{'} \mathbf{c}$를 얻는다. 

 

STEP2. 각각의 mask sequence를 2D mask로 reshape하여 $\mathbf{s}_{\text{mask}} \in \mathbb{R}^{H / P \times W / P \times K}$를 얻는다. 

 

STEP3. $\mathbf{s}_{\text{mask}}$에 bilinear upsampling을 통해 원본영상 크기 $\mathbf{s} \in \mathbb{R}^{H \times W \times K}$를 얻는다. 

 

STEP4. Softmax를 class dimension에 적용하여 final segmentation map을 얻는다. 

 

4) Model Variants

 

 

Experiment Results

1) Ablation Study

1-1) Regularization

 

 

1-2) Transformer Size and Patch Size

 

 

1-3) Decoder Variants

 

2) Quantitative Results

 

 

3) Qualitative Results