[P-stage 1] Swin Transformer

Transformer를 사용한 모델도 사용해보자는 의견이 있어 정리하고 사용해보기 위한 정리글.

ViT (Vision Transformer)

NLP 분야에서 한 획을 그은 (GPT 모델에도 사용되는) Transformer를 image classification 분야에 맞게 변형하며 CNN을 사용하지 않도록 만든 모델.

사실상 이 모델의 등장 이후 CNN을 사용한 모델의 종말을 예고했고, Transformer가 NLP뿐 아닌 CV까지 뻗치며 인공지능의 근본이 되었다.

1. 이미지를 여러 개의 patch로 자른 후 각 패치별 embedding demension (16 x 16 x 3)으로 만든다. 이를 통해 기존 Transformer 모델에서의 시퀀스 데이터와 같이 만든다.

2. 각 패치에 대한 flatten 작업을 통해 벡터화하고, Linear 연산을 통한 Embedding을 진행한다.

3. Embedding 결과에 클래스를 예측하는 클래스 토큰을 추가한다.

4. 클래스 토큰이 추가된 입력값에 Positional Embedding을 더해 위치에 따른 정보를 추가한다.

이 과정을 통해 ViT의 입력값을 생성한다. 이후 Transforemr Encoder를 L번 반복하면 입력값과 동일한 크기의 출력값을 얻을 수 있고, 여기에 대한 MLP 연산을 통해 최종 class 예측이 가능하다.

 

하지만 이 ViT의 단점으로는 엄청나게 많은 데이터가 필요하다. 따라서 엄청나게 많은 파라미터가 필요하고, 연산량 역시 많다. 이를 개선한 점이 Swin Transformer이다.

 

Swin Transformer

ViT에서 모든 patch에 대한 self-attention 연산이 지나치게 많이 드는 점을 개선하고자 patch를 window로 나누어 해당 윈도우 내에서만 self-attention을 수행하고 이 window를 shift해 다시 연산을 하는 구조를 제안한다.

입력 이미지 사이즈 224x224에 대해 ViT에서는 16 x 16 픽셀을 만들어 총 196개의 패치에 대해 self-attention을 수행한다.

Swin Transformers에서는 FPN 구조처럼 patch 4x4 사이즈에서 이 patch에 대한 merge 연산을 더해가며 연산을 수행한다. 여기서 patch 내에서 빨간 선으로 나누어진 부분을 window라고 부르며, 여기에 대해 window 내의 patch에 대해서만 self-attention을 수행한다.

 

전체적인 Swin Transformer의 구조. 크게 Patch Partition, Linear Embedding, Swin Transformer, Patch Merging으로 구분이 되며 4개의 stage로 이루어져있다.

Swin Transformer block은 두 개의 encoder로 구성되어 있으며 일반적인 Multi head self-attention이 아닌 W-MSA, SW-MSA로 이루어져있다. 

 

Patch Partitioning

ViT와 같은 patch 분리 모듈을 통해 입력 RGB 이미지를 겹치지 않는 patch들로 나눈다. 각 patch는 토큰의 역할을 하며, 논문에서는 4x4크기의 3채널 48feature를 사용한다.

이후 Linear Layer를 통해 H/4 x W/4 x 48 텐서를 H/4 x W/4 x C 텐서로 변환한다. 이후 Swin Transforemr Block에서는 W-MSA와 SW-MSA를 사용해 GELU 활성화 함수를 사용한다.

 

Patch Merging

맨 처음의 4x4의 patch에 대해 merge 연산을 통해 더 넓은 부분을 한번에 보는 과정이다. 인접한 2x2개의 patch를 concat 한 후 Linear Layer를 통해 2C의 사이즈로 맞춘다. 따라서 채널이 2배가 되며 해상도가 2x2배 줄어든다.

 

Shifted Window based SelfAttention

W-MSA와 SW-MSA는 기존 ViT의 제곱되는 연산량을 해결하기 위한 방법이다. 

가장 작은 patch가 4x4 사이즈이며 Window는 M x M개의 patch로 구성한다. 논문에서는 7x7을 사용한다.

이미지에는 이 윈도우가 h x w개 있어 해상도가 4Mh x 4Mw이다.

이미지 해상도 224x224 기준 첫 layer의 patch는 4x4, 이것을 7x7개 모은 window가 8x8개 존재한다. 이경우

W-MSA의 연산량은 다음과 같이 줄어든다.

Window 방식에서 경계의 patch들은 인접한 상태에서도 attention 연산을 수행할 수 없다. 따라서 Window를 shift해서 움직이며 계산을 수행한다. 이 때 shift하는 cyclic은 M/2이다.

위와 같은 과정을 따르면 이전 layer의 입력이 들어오면 W-MSA를 수행해 MLP를 통과시킨다. 이후 SW-MSA를 수행한 후, 다시 MLP를 통과한다. 그러면 Swin Transformer block의 output이 다음 layer로 전달된다. 

이 과정을 정리하면

1. 이전 layer의 입력이 들어온다

2. W-MSA를 수행한다.

3. MLP에 통과시킨다.

4. SW-MSA를 수행한다. 이후 Reverse cyclic shift를 수행한다.

5. MLP를 통과시켜 output을 도출한다.

6. 이를 다음 layer로 전달한다.

 

기존의 ViT와 비교해 Swin Transformer의 장점은 기존의 ViT와 달리 계산량을 많이 줄이며 다양한 이미지를 처리할 수 있다. 또한 여러 비전 task에서의 backbone으로 활용할 수 있다.

 

 

 

 

Reference

https://gaussian37.github.io/dl-concept-vit/

Vision Transformer (AN IMAGE IS WORTH 16X16 WORDS, TRANSFORMERS FOR IMAGE RECOGNITION AT SCALE)

https://visionhong.tistory.com/31

[논문리뷰] Swin Transformer

https://greeksharifa.github.io/computer%20vision/2021/12/14/Swin-Transformer/

Python, Machine & Deep Learning