[논문 리뷰]Rethinking Semantic Segmentation from a Sequence-to-Sequence Perspective with Transformers
들어가며
이번에 리뷰할 논문은 2021년 CVPR에서 발표된 Rethinking semantic segmentation from a sequence-to-sequence perspective with transformers 입니다.
제목에 나와있듯이 논문에서는 transformer 구조를 사용하여 semantic segmentation 모델을 구현하였습니다.
들어가기에 앞서, semantic segmentation과 transformer가 무엇인지 간단히 정리하고 넘어가겠습니다.
- Semantic segmentation은 이미지에서 객체를 픽셀단위로 분류하는 작업입니다.
간단히 말해서 아래 그림처럼 입력 이미지가 들어왔을 때, 픽셀별로 객체에 해당하는 클래스로 라벨링합니다.
- Transformer는 2017년 “Attention is all you need” 논문에서 제안된 딥러닝 모델로 attention 구조만을 사용했습니다.
Seq2seq 모델로 처음에는 자연어 처리 분야에서만 사용됐지만, 2020년에 ViT(Vision transformer)가 나오면서 이미지 처리 분야에서도 널리 사용되고 있습니다.
Transformer에 대한 내용은 잘 정리된 글들이 많으니, 논문이나 다른 블로그를 참고 바랍니다.
Introduction
2015년 CVPR에 “Fully convolutional networks for semantic segmentation.” 논문이 나온 후로, semantic segmentation 모델은 주로 FCN(Fully convolutional network)을 기반으로 설계되어왔습니다.
FCN 구조를 기반으로 한 segmentation 모델은 CNN으로만 이루어진 인코더와 디코더로 구성되어 있습니다.
Input 이미지가 들어오면, 인코더에서는 feature map을 downsampling하면서 semantic segmentation을 위한 특징들을 학습하고, 디코더에서는 학습된 feature를 다시 원본 이미지 크기로 upsampling 시켜 픽셀별 classification을 수행할 수 있게끔합니다.
FCN 기반의 모델은 layer를 지날수록 feature map의 resolution이 낮아지면서 receptive field가 커지긴 하지만, convolution filter만 사용함에 따른 이미지의 전체적인 global information(Long-range dependency information)을 학습하지 못한다는 한계점이 존재합니다.
이러한 한계점을 극복하기 위해서 이전 논문들에선 convolution 연산을 기존과 다르게 하거나, FCN에 attention 구조를 혼합하여 사용해왔습니다.
하지만, FCN의 제한된 receptive field는 long-range dependency 문제를 완벽히 해결하지 못하고, attention 구조는 계산 비용 문제로 FCN의 상위 layer에서밖에 사용하지 않았습니다.
이러한 한계점을 해결하기 위해, 논문에서는 transformer를 사용하여 semantic segmentation을 해결하는 SEgmentation Transformer(SETR)을 제안합니다.
인코더로 transformer를 사용함으로써 long-range dependency 문제를 해결하고, FCN 구조에서 벗어나기 위해 인코더에서 동일한 크기의 feature map을 사용하였습니다.
Contributions
논문에서 말하는 기여는 다음과 같습니다.
- 기존의 FCN 구조를 벗어난 모델을 제안하고, seq2seq 학습 관점에 semantic segmentation 문제를 재구성한다.
- 이미지를 패치로 잘라서 sequence 형태로 입력을 받는 transformer 인코더를 제안한다.
- 서로 다른 디자인을 갖는 3가지 디코더를 제안한다.
- ADE20K, Pascal Context, Cityscapes 데이터셋에서 SOTA를 달성했다.
Method
Encoder
인코더 부분은 ViT와 거의 다를게 없기 때문에 “An image is worth 16x16 words: Transformers for image recognition at scale” 논문을 먼저 읽고 오시는 걸 추천드립니다.
먼저, H x W x 3 차원을 갖는 이미지가 주어지면 transformer에 들어가기 전에 sequence 형태로 바꿔주는 embedding 과정을 거치게 됩니다.
L개의 patch로 입력 이미지를 분할하는데, 일반적으로 기존의 semantic segmentation 인코더가 입력 이미지에 대한 최종 feature map을 H/16 x W/16까지 down sampling 하는 것을 고려해서 L = H/16 x W/16로 설정합니다.
하나의 이미지 patch는 16 x 16 x 3 차원을 갖게 되고, flatten 하게 되면 C = 16 x 16 x 3개의 값을 갖는 하나의 sequence로 생각할 수 있습니다.
결과적으로 H x W x 3 차원을 갖는 입력 이미지는, 하나당 C개 값을 갖는 L개의 sequence로 나타낼 수 있습니다.
마지막으로 각 sequence에 해당되는 patch의 위치 정보를 더해주는 positional embedding을 진행합니다.
위 과정을 모두 마치면 transformer의 입력으로 L x C 차원의 embedding sequence를 얻게 됩니다.
SETR은 24개의 layer로 구성되고, 각 layer는 layer normalization, MSA(Multi-head self-attention), MLP(Multilayer perceptron)가 residual 구조로 이루어져있습니다.
그리고 각 layer의 입력과 출력은 L x C 차원으로 고정되어 있어, 기존의 semantic segmentation 모델과 다르게 down sampling이 일어나지 않습니다.
모든 layer에서 입력에 layer normalization을 적용하고 self-attention을 진행합니다.
l번째 layer의 key, query, value 값은 각각에 해당하는 weight W와 이전 layer의 출력을 곱해서 구할 수 있습니다.
W는 모두 C x d 차원을 갖고, 계산된 query, key, value는 모두 L x d 차원을 갖는 것을 확인할 수 있습니다.
key와 query를 곱하여 L x L 차원의 값을 갖고, 다시 value와 곱하여 L x d 차원의 값을 갖습니다.
입력과 residual connection을 해주는데 입력의 차원이 L x C 이므로, 논문에 나와있지 않지만 d = C로 설정했다고 생각할 수 있습니다.
MSA는 m개의 head에서 각각 self-attention을 나눠서 진행한 뒤 합쳐주며 계산량을 줄여주는 역할을 합니다.
MSA 과정에서 key, query, value의 weight W는 모두 C x d/m 차원을 갖고, 각 head의 self-attention 결과를 concatenate 해서 L x C 차원의 결과를 출력합니다.
MLP는 여러 head에서 나온 결과를 합쳐주는 fully connected layer라고 간단히 생각하면 될 것 같습니다.
수식엔 나와있지 않지만 MSA, MLP 전에 layer normalization을 취해주고, residual connection에는 layer normalization 취하기 전 값을 사용합니다.
Decoder
이제, pixel-level의 segmentation을 진행하기 위해 Decoder를 거치게 됩니다.
Decoder를 거치면서 encoder feature, 즉 encoder의 최종 결과물인 Z를 평가하게 됩니다.
이 과정에서 encoder feature인 Z의 shape이 HW/256 X C 에서 H/16 * W/16 * C로 변환됩니다.
저자는 총 세 가지의 decoder를 제안하였습니다.
1.Naive upsampling (Naive) — SETR-Naive
첫 번째로, Navis upsampling입니다.
Transformer feature인 Z^Le category number로 project한 후, 2-layer network를 적용합니다.
network 구성은 하단과 같습니다.
1x1 conv + sync batc norm (w/ ReLU) + 1x1 conv
그 후 bilinearly upsample하여 full image resolution으로 확대합니다.
이 과정에서 pixel-wise cross entropy loss를 사용합니다.
2.Progressive UPsampling (PUP) — SETR-PUP
두 번째 decoding 방식은 Progressive UPsampling (PUP) 방식입니다.
noisy prediction을 유발하는 One-step upscaling 시, noisy prediction이 유발되는 문제를 제거하기 위해 progressive upsampling 전략을 이용하였습니다.
adversarial effect를 최대한 완화시키기 위해 Upsampling을 한 번에 2배만 가능하도록 제한하였고, 이 때문에 Z^Le를 H/16*W/16 해상도로 원복하기 위해서 총 4번의 operation이 필요하게 됩니다.
3.Multi-Level feature Aggregation (MLA) — SETR-MLA
마지막 decoding 방식은 Multi-Level feature Aggregation(MLA)입니다.
위 그림을 보면 언뜻 Feature pyramid network와 비슷해 보이지만, MLA과 차이점은 각 layer의 feature representation Z^l은 피라미드 shape없이 same resolution을 가진다는 점입니다.
Transformer feature representation인 Z^m을 다음과 같이 선별하여 사용합니다.
위 그림처럼 m=4, Le=24인 경우 Z⁶, Z¹², Z¹⁸, Z²⁴를 사용하게 됩니다.
각 stream에서 Z^l는 HW/256C에서 H/16W/16 *C으로 reshape 일어납니다.
3-layer (kernel size 1 x 1, 3 x 3, 3 x 3) network를 적용하여, 첫번째와 세번째 레이어의 특징을 절반으로 줄이고, 세번째 레이어 이후 bilinear 연산을 통해 spatial 해상도를 4배 높힙니다.
서로 다른 stream들의 상호작용을 위해, element-wise addition인 top-down aggregation를 적용한 후, 추가로 추가 3x3 conv를 적용합니다.
세번째 레이어 이후, channel-wise concatenation하여 모든 stream에서 융합된 feature를 얻고, 그 후 전체 해상도로 bilinearly 4배 upsampling 하게 됩니다.
Experiment
자 그럼 SETR의 성능을 살펴볼까요?
위 그림은 Cityscapes로 진행된 정성적 실험결과입니다.
왼쪽 그림은 dilated FCN의 결과이고, 오른쪽 그림이 SETR의 결과입니다.
그림에서 알 수 있듯이, segmentation된 부분 안에 이물질이 들어간 듯해 보입니다.
Segmentation이 깔끔하게 진행되지 않았다는 의미죠.
하지만 SETR 모델의 결과인 오른쪽 그림은 왼쪽 그림에 비해 깔끔하고 정확하게 Segmentation된 것을 볼 수 있습니다.
위 표는 Cityscapes를 이용하여 State-of-the-art(SOTA)와 비교한 성능표입니다.
여기서 오늘 소개해드린 SETR-PUP 모드가 제일 좋은 성능을 보여줌을 알 수 있습니다.
본 논문에서는 Cityscapes외에도 Abliation studies를 포함하여 ADE20K, PASCAL Context로 정성적, 정량적 실험한 결과가 첨부되어 있습니다.
더 많은 실험 결과는 페이퍼를 참고해주세요.
Conclusion
Semantic Segmentation을 Sequence to sequence 분류 문제로 바꾸어 해결한 SETR에 대해 살펴봤습니다.
FCN을 transformer로 대체하여 기존의 FCN의 receptive field를 변형하던 component level에서 architectural level로 모델에 큰 변화를 주었는데요.
기존의 FCN 부분을 모두 덜어내어 transformer 인코더를 사용하고, 서로 다른 디자인을 갖는 3가지 디코더를 적제적소에 배치하여, 실험에서도 좋은 성능을 보여주었습니다.
Klleon에서도 얼굴 변환을 하기 전, face segmentation을 하기 위해 semantic segmentation 기술을 사용하는데요, 기존의 Semantic Segmentation보다 더 빠른 속도와 높은 성능을 보장하기 위한 연구가 진행되고 있습니다.
그럼 다음 포스트에서도 오늘처럼 기발한 아이디어를 가진 논문 리뷰로 다시 찾아뵙도록 하겠습니다.
Reference
[1] Zheng, Sixiao, et al. “Rethinking semantic segmentation from a sequence-to-sequence perspective with transformers.” Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2021.
[2] Long, Jonathan, Evan Shelhamer, and Trevor Darrell. “Fully convolutional networks for semantic segmentation.” Proceedings of the IEEE conference on computer vision and pattern recognition. 2015.
[3] Vaswani, Ashish, et al. “Attention is all you need.” Advances in neural information processing systems 30 (2017).
[4] Dosovitskiy, Alexey, et al. “An image is worth 16x16 words: Transformers for image recognition at scale.” arXiv preprint arXiv:2010.11929 (2020).
