Code Review: Vision Transformer (Ref: Google research github)

Vision Transformer (ViT) 모델 구조를 코드와 함께 리뷰해보려합니다.

 

Paper: https://arxiv.org/abs/2010.11929

Offifical github: https://github.com/google-research/vision_transformer

GitHub - google-research/vision_transformer

 

Vision Transformer (ViT) 모델 구조를 코드를 함께 리뷰해보려합니다. 정확한 리뷰를 위하여 official github에 업로드된 코드를 가지고 리뷰를 할것이며 ViT 모델은 vit_jax/model.py 파일에 작성되어있습니다.

 

model.py는 5개의 class로 구성되어 있으며 class 목록은 아래와 같습니다.

 - AddPositionEmbs

 - MlpBlock

 - Encoder1DBlock

 - Encoder

 - VisionTransformer

 

5개의 class를 한눈에 쉽게 이해하기 위하여 조직도로 그려보았습니다.

 

즉, 5개의 class를 이해하면 Vision Transformer 모델의 전체 구조를 파악할 수 있는 것입니다. 모델의 구조를 좀 더 쉽게 이해하기 위해 모델에 사용되는 변수들에 값을 부여하도록 하겠습니다. 논문에 작성되어 있는 ViT-Base, ViT-Large, ViT-Huge 중, ViT-Large(ViT-L)를 사용하도록 할 것이며 정보는 아래와 같습니다. 

patch size는 16을 사용할 것이며 최종 사용하는 모델은 ViT-L/16 입니다.

 

그럼 파악해보도록 하겠습니다. 

 

 

 VisionTransformer 

>> Step1 

논문에 작성되어 있는 모델을 설명하는 그림을 보면 이미지 하나가 patch로 쪼개진 다음 Linear Projection과 Flatten Patches 과정을 통과한다. 

논문을 보면 2D 이미지 (HxWxC)를 2D 패치 (P^2, C)로 flatten 시켜준다고 작성되어있습니다.

위 논문 내용에 매칭되는 코드 부분은 아래와 같습니다. 

• 여기서 x는 patch로 자르기 전 이미지 묶음입니다. 즉, batch단위로 묶인 이미지 집합들인것입니다. 숫자로 따라가보기 위해 임의로 하나의 이미지 크기를 (224, 224, 3)로 batch를 512로 정의하겠습니다. 그렇다면 Line 255의 x.shape은 (512, 224, 224, 3)인것입니다. 그렇다면 n=512, h=224, w=224, c=3인것입니다. 
• 그 다음 Line 258을 보면 갑자기 Convolution이 나오게 됩니다. 논문에는 이미지를 patch로 reshape 한다고만 명시되어있습니다. (논문에는 간략히 명시만 되어있기에 논문에서 제안하는 모델을 명확하게 파악하려면 논문과 코드를 같이 보는 것은 중요하다는 점 다시 알려드리며 마저 리뷰하도록 하겠습니다.) 여기서는 reshape하는 과정에 convolution이 포함되어 있습니다. 그런데 코드를 보면 kernel size와 stride size가 patch size로 동일하게 정의되어 있습니다. 이 말의 의미는 feature를 추출하는 영역이 겹치지 않다는 뜻입니다. 왜냐하면 특징을 추출하고자 하는 영역만큼 이동한 다음 특징을 추출하기 때문입니다. 그렇게 되면 해당 convolution을 통과한 output 크기는 어떻게 될까요? 이미지의 가로 크기를 patch 가로 크기로 나눈 값이 output의 가로 크기가 되고, 이미지의 세로 크기를 patch 세로 크기로 나누 값이 output의 세로 크기가 됩니다. 그리고 코드에서 output의 차원은 self.hidden_size만큼 되게 됩니다. (위 코드에는 명시가 안되어있지만 논문에 self.hidden_size는 1024라고 명시되어있습니다.)
• 그렇다면 Line 269에서 convolution을 통과한 x의 shape은 (512, 14, 14, 1024)가 됩니다. 위 코드의 목적은 2차원으로 reshape하기 위함입니다. 그러므로 Line 270에서 x를 reshape하게 되면 x의 shape은 (512, 196, 1024)가 됩니다.(Line 270의 hxw는 patch의 개수로 196입니다. 논문에는 N=HW/P^2라고 명시되어있는데 논문에서의 H와W는 처음 이미지의 height와 width므로 계산하면 196이 나오게됩니다.)

 

>> Step2 

그 다음 단계는 Patch에 Position Embedding 과정을 거치는 것을 확인할 수 있습니다.

 

논문을 보면 Position embedding이란 patch embedding에 positional information을 추가해주는 과정을 의미한다고 합니다.

 

위 논문 내용에 매칭되는 코드 부분은 아래와 같습니다.

• 0으로 초기화된 (1, 1, 1024) 차원의 배열을 만든 후, 배열을 붙이는 tile함수를 이용하여 (512, 1, 1024) 차원의 배열로 만들어 줍니다. 후에 concatenate함수를 이용하면 x의 shape은 (512, 197, 1024)가 됩니다.
• 이로써 196개 patch embeddings에 하나의 학습가능한 1D position embedding이 추가된 x가 완성되었습니다.
• 후에 각 patch에 position embedding을 추가해줘야 step2의 그림에 나와있는 부분이 완성됩니다. 해당 부분은 클래스 AddPositionEmbs에서 완성되며 해당 클래스는 클래스 Encoder가 불러오게 됩니다. AddPositionEmbs 클래스에 작성되어있는 코드의 양이 적기에 바로 이어서 설명하도록 하겠습니다.

 

 AddPositionEmbs 

앞서 설명한 각 patch에 position embedding을 추가하는 코드 부분은 아래와 같습니다. 

• inputs의 shape은 (512, 197, 1024)이기에 pos_emb_shape은 (1, 197, 1024)가 됩니다.
• inputs과 pe를 더해줌으로써 각 patch에 position embedding이 추가되고, 해당 값은 반환됩니다.

 

 Encoder 

Encoder 클래스내 상단부분을 확인해보면 AddPositionEmbs 클래스를 이용하여 각 patch에 position embedding을 추가한 후, Dropout 함수를 거칩니다. 이때, self.dropout_rate 값은 논문에 0.1이라고 명시되어있습니다. 

Encoder 클래스를 처음에 접했을 때는 Vison Transformer model내에 있는 Transformer Encoder의 구조가 작성되어있을 것이라고 파악했는데, position embedding 추가하는 코드까지 해당 부분에 작성되어있습니다. 그렇기에 해당 클래스를 Transformer Encoder의 설명만이 작성되어있다기보다는 Transformer Encoder에 입력되는 input 처리 과정과 Transofmrer Encoder 구조가 작성되어있다고 보는게 맞을 것 같습니다.

Norm > Multi-Head Attention > Norm > MLP 과정을 self.num_layers 만큼 거치게 됩니다. 거쳐야 할 4개의 과정을 Encoder1DBlock 클래스로 정의한 것을 알 수 있습니다. 이 후, Layer Normalization를 통과하게 되면 Transformer Encoder가 끝나게 됩니다. 

Vision Transformer 모델의 핵심 부분인데, 짧게 느껴지는 이유는 핵심인 부분을 Encoder1DBlock 클래스로 처리해서이기에 바로 Encoder1DBlock 클래스에 대해서 알아보도록 하겠습니다.

 

 Encoder1DBlock 

먼저 Layer Normalization을 통과한 후, MultiHeadDotProductAttention를 통과하게 됩니다. 논문에 MultiHeadDotProductAttention는 Multiheaded self-attention 라고 작성되어있습니다. Multi-Headed Attention(MHA)에 대해서 간략히 설명드리자면 Query, Key, Value가 존재합니다. Query는 영향을 받을 주체이고 Key는 영향을 끼칠 주체 입니다. 즉 2개의 입력정보 주체가 존재하는 것입니다. 이것은 Multi-Headed Attention일 경우이며, Multi-Headed Self Attention(MSA)일 경우에는 입력정보 주체가 1개로 변하게 됩니다. 즉, Query와 Key가 동일해지는 것입니다. MSA와 MHA는 입력정보 주체가 다른 것일 뿐 모델 구조는 동일합니다. MHA와 MSA가 기능적으로 훌륭한 성능을 냄에 따라 용이한 사용을 위해 기본적으로 함수가 제공되었고, 해당 github에서는 제공하는 함수를 사용하였습니다.

이후, 마지막으로 self.dropout_rate로 rate값이 설정된 Dropout을 통과하고 해당 값을 초기의 Layer Normalization의 입력과 더합니다.

 

그 다음, Layer Normalization을 통과한 후, MlpBlock 클래스를 통과합니다. 해당 값과 두 번째 Layer Normalization의 입력과 더합니다.

Multiheaded Self Attention(MSA)와 MLP block은 구조도 다르지만 본 논문에서 주장하는 부분은 MSA는 전역적으로 특징을 추출할 수 있는 반면, MLP block은 지역적으로 특징을 추출한다는 것입니다. 즉, 전역적 특징 추출과 지역적 특징 추출을 동시에 사용함으로써 어느 한쪽으로 치우쳐지지않고 두 개의 장점을 다 가져가기 위함입니다.

 

 MlpBlock 

MLP는 Multi-Layer Perceptron의 줄임말로, 여러개의 perceptrion을 여러 층으로 쌓은 다층신경망 구조입니다. 즉, 입력층과 출력층 사이에 최소 하나 이상의 은닉층을 갖고 있는 구조입니다. Dense Layer > GeLU > Dropout > Dense Layer > Dropout 과정을 거치게 됩니다. 이 중, 생소한 부분인 GeLU(Gaussian Error Linear Unit) activation function에 관련된 설명은 필자가 작성한 다른 게시물을 참고하시기 바랍니다.(https://s5unnyjjj.tistory.com/78)

---

>> 위 내용과 그림은 필자가 직접 작성한 내용입니다.

>> 부족한 점이 많을 수 있기에 잘못된 내용이나 궁금한 사항이 있으면 댓글 달아주시기 바랍니다.

>> 긴 글 읽어주셔서 감사합니다.