안녕하세요. 지난 포스팅의 [IS2D] SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformer (NIPS2021)에서는 Transformer 기반의 의미론적 분할 모델인 SegFormer에 대해서 소개하였습니다. 오늘은 그보다도 Transformer 기반 의미론적 분할 모델의 원조격과 같은 느낌인 SETR에 대해서 소개하도록 하겠습니다 (Object Detection 모델인 DETR이란 비슷하네요!). SETR은 SegFormer에서 실험비교에서도 나왔었죠.
Background
기본적으로 의미론적 영상 분할 (Semantic Segmentation)은 영상 내의 각 픽셀에 대해 어떤 카테고리 레이블에 속하는 지 예측하는 dense prediction 문제라고 볼 수 있습니다. 이를 통해 사용자는 전체 영상 내에 완벽한 이해 (complete understanding)를 제공해주죠. 이러한 문제는 자율주행, 이상탐지, 그리고 의료영상 등과 같은 정말 다양한 도메인에서 필수 기술로서 자리매김하고 있습니다.
이를 위해 제안된 다양한 모델들이 있습니다. 대표적으로 CNN-based 모델들은 inductive bias를 기반으로 설계되어 상대적으로 적은 양의 데이터로도 충분히 좋은 성능을 얻을 수 있다는 장점은 있으나 global information을 활용하지 못합니다. 대표적으로 FCN (CVPR2016)과 DeepLab Series (ICLR2015, TPAMI2017, arxiv2017, ECCV2018) 등이 있습니다. 하지만 이러한 모델들 모두 FCN과 유사한 encoder-decoder 형태의 모델에서 벗어나지 못하고 있습니다.
본 논문에서는 이러한 한계점에 유의하며 오직 Transformer로만 구성된 의미론적 분할 모델을 설계해보고자 합니다. 이렇게 제안된 SEgmentation TRansformer (SETR)은 stacked convolution encoder를 Transformer로 변경하였죠. 핵심적으로 SETR은 patch embedding을 통해 입력 영상을 패치의 시퀀스로서 해석하여 ViT와 동일한 컨셉으로 토큰을 추출하는 것으로 시작합니다. 그리고 결정적인 feature representation을 학습하기 위해 패치 시퀀스를 global self-attention을 이용해 global information을 활용할 수 있게 만들어줍니다. 이 과정에서 downsampling이 적용되지 않아 upsampling이 필요가 없습니다. 이러한 SETR은 ADE20K, PASCAL Context, CityScapes와 같은 다양한 벤치마크 데이터셋에서 높은 성능을 달성하였습니다.
SETR
1) FCN-based Semantic Segmentation
본격적으로 시작하기에 앞서 FCN의 특징들과 동작원리를 다음과 같이 정리해볼 수 있습니다.
- FCN encoder는 여러 개의 convolution block이 쌓여 구성되어 있다.
- 첫번째 계층은 $H \times W \times 3$ 크기의 입력 영상을 받는다.
- 두번째 계층부터 마지막 계층까지 3차원 $h \times w \times d$의 형태로 다음 계층으로 전달된다.
- 깊은 계층의 출력은 얕은 계층의 특징맵의 공간정보를 기반으로 연결되어 receptive field라고 하는 영역에서 연산이 고려된다.
이러한 특징들과 동작원리로 인해 FCN은 근본적으로 오직 가장 깊은 계층에서만 가장 큰 receptive field를 가지게 되어 long-range dependency를 모델링할 수 있습니다. 이는 다양한 계층에서 long-range dependency를 모델링할 수 없는 것으로 이해할 수 있기 때문에 FCN의 가장 큰 단점이라고 볼 수 있죠.
2) Overall Architecture
SETR의 전체적인 모델 구조는 그림 1 (a)과 같습니다. 언뜻보았을 때 ViT와 별 차이점이 느껴지지 않습니다. 다만, 영상 분류를 위주로 설계된 ViT와는 다르게 SETR은 의미론적 영상 분할을 목표로 설계되었기 때문에 Transformer 인코더 이후에 디코더가 존재합니다. 그리고 그림 1 (b)와 (c)는 각각 해당 디코더에서 사용할 수 있는 서로 다른 구조 입니다. (b) 같은 경우에는 FCN 또는 UNet 구조에서 많이 사용하는 upsampling-conv 연산이 반복되는 디코더로 Progressive Upsample이라고 하죠. 그리고 (c) 같은 경우에는 다중 계층 정보를 활용하여 하나로 결합한 뒤 최종 출력을 내는 Multi-Level feature Aggregation 구조로 본 논문에서 채택한 디코더 구조입니다.
2-1) Image-to-Sequence
SETR은 기본적으로 ViT에서 사용한 Transformer를 거의 그대로 차용합니다. 따라서, 3D 입력 영상을 2D 시퀀스로 매핑하는 과정이 필수적이죠. 즉, Transformer 인코더는 2D 시퀀스 $Z \in \mathbb{R}^{L \times C}$를 입력으로 받으며 여기서 $L$은 토큰의 개수, $C$는 임베딩 벡터의 길이를 의미합니다. 따라서, 저희는 입력 영상 $x \in \mathbb{R}^{H \times W \times 3}$을 시퀀스 $Z \in \mathbb{R}^{L \times C}$로 image sequentialization을 수행해야하죠. 다만, 토큰의 개수는 곧 연산량에 직결되기 때문에 몇 개의 토큰으로 나눌 지는 매우 중요한 요소입니다. 예를 들어, 픽셀 별로 토큰을 나눈다라고 가정했을 때 $224 \times 224$ 크기의 입력 영상에서는 50176개의 토큰이 나오므로 엄청난 양의 연산량이 요구됩니다. 따라서, ViT에서도 픽셀 단위로 아닌 패치 단위로 토큰을 나누듯이 SETR에서도 패치 단위로 나누어줍니다.
일반적으로 영상 분할 모델에서는 특징 맵을 $\frac{H}{16} \times \frac{W}{16} \times C$까지 줄여서 진행하는 경우가 많습니다. SETR에서는 이러한 특성에 주목하며 각 패치의 크기를 $16 \times 16$, 즉 전체 토큰의 개수를 $\frac{H}{16} \times \frac{W}{16} = \frac{HW}{256}$개로 잡고 진행합니다. 이렇게 되면 이전에 픽셀 단위의 토큰에서는 50176개였지만 이번에는 196개로 큰 감소를 하게 되죠.
이제 각 패치 토큰들은 일반적인 ViT에서 하던 방법과 동일하게 flatten (vectorization)이 수행되고 linear projection을 통해 $C$개의 채널을 가지도록 매핑됩니다. 마지막으로 position 정보를 주입해주기 위해서 positional embedding vector를 각 임베팅 토큰 별로 더해주게 됩니다.
2-2) Transformer
다음으로 총 $L_{e}$개의 Transformer layer로 구성된 인코더를 통과합니다. 여기서 사용하는 Transformer layer는 여러분들도 다 아시다싶이 MHSA와 MLP로 구성되어 있으므로 자세한 설명을 생략하도록 하겠습니다.
2-3) Decoder Design
마지막으로 pixel-wise prediction을 위한 decoder를 설계해야합니다. 본 논문에서는 총 3가지의 설계를 제시합니다.
- Naive Upsampling (Naive): 가장 단순한 방식으로 Transformer 인코더를 통과해나온 feature에 linear projection을 통해 전체 카테고리 개수만큼으로 채널 개수를 축소시킵니다. 이 과정은 $1 \times 1$ conv + sync batch normalization + ReLU + $1 \times 1$ conv로 구성되어 있죠. 마지막으로 Bilinear interpolation을 통해 입력 영상과 동일한 크기를 가질 수 있도록 만들어줍니다.
- Progressive Upsample (PUP): 일반적으로 FCN 구조에서 많이 사용하는 방식으로 conv layer와 upsampling을 번갈아가며 진행하는 방식입니다. 이 과정에서 upsampling은 최대 $2$배만큼만 늘어나도록 하여 극단적인 정보손실을 방지하죠. 따라서, 입력 영상만큼의 해상도로 다시 복원되기 위해서는 총 4번의 연산을 반복해야합니다.
- Multi-Level feature Aggregation (MLA): feature representation $\{ Z^{m} \}$ 중에서 $M$개의 Transformer 인코더 계층에서 나온 feature map을 균등하게 선택한 뒤 top-down 방식으로 aggregation을 수행합니다. 그리고 conv+conv를 거친후 하나의 feature로 concat한 뒤 다시 한번 conv를 거쳐 bilinear interpolation을 통해 입력 영상과 동일한 크기를 가질 수 있도록 만들어줍니다.
2-4) Model Variants
Experiment Results
1) Ablation Study
1-1) Model Variants
1-2) Model Pre-training
2) ADE20K & PASCAL Context
3) CityScapes
4) Feature Visualization
