Vision Transformer(ViT) 공부 및 간단 구현 예제

Vision Transformer (ViT)는 이미지 처리 분야에서 사용되는 트랜스포머 아키텍처에 기반한 모델.

2020년에 Google Research에서 "AN IMAGE IS WORTH 16X16 WORDS: TRANSFORMERS FOR IMAGE RECOGNITION AT SCALE"이라는 논문을 통해 소개되었습니다. ViT는 주로 자연어 처리에서 매우 성공적이었던 트랜스포머를 이미지 인식에 적용하는 혁신적인 접근법을 제시합니다.

ViT의 핵심 아이디어 및 구조는 다음과 같습니다:

1. 이미지 패치: 이미지를 여러 개의 고정 크기 패치로 나눕니다. 예를 들어, 224x224 크기의 이미지를 16x16 크기의 패치로 나누면 14x14=196개의 패치가 생성됩니다.
2. 패치 임베딩: 각 패치를 벡터로 평탄화(flatten)하고, 이것에 임베딩 레이어를 적용하여 길이가 D인 벡터로 변환합니다. 이 D-차원 벡터는 패치의 정보를 나타냅니다.
3. 포지션 임베딩: 이미지 내의 각 패치 위치에 대한 정보를 추가하기 위해 포지션 임베딩을 사용합니다. 이 임베딩은 각 패치 벡터에 더해져 위치 정보를 제공합니다.
4. 트랜스포머 인코더: 포지션 임베딩이 추가된 패치 임베딩들은 트랜스포머 인코더로 전달됩니다. 여기서 패치들 간의 관계와 컨텍스트를 학습합니다.
5. 분류 헤더: 트랜스포머의 출력은 풀링 레이어나 분류 레이어로 전달되어 최종 예측을 수행합니다.

ViT는 충분한 양의 데이터와 컴퓨팅 리소스가 있을 때 특히 성능이 뛰어납니다. 초기 학습 시 큰 데이터셋에서 사전 학습을 진행한 후, 특정 작업에 맞게 미세 조정(fine-tuning)을 하는 방법이 흔히 사용됩니다.

ViT의 등장 이후, 다양한 트랜스포머 기반의 이미지 처리 아키텍처가 연구되고 있으며, 그 중 일부는 이미지 처리에 있어 기존의 CNN 아키텍처를 능가하는 성능을 보이기도 합니다.

---

ViT 간단히 구현한 예제 코드입니다.

1. 라이브러리 가져오기

import torch
import torch.nn as nn

• torch: 파이토치라는 딥러닝 프레임워크를 사용하기 위한 라이브러리입니다.
• torch.nn: 딥러닝 모델을 만들기 위한 기본 구성 요소들을 포함하고 있습니다.

---

2. 이미지를 패치로 나누는 모듈

class PatchEmbedding(nn.Module):
    def __init__(self, in_channels: int = 3, patch_size: int = 16, emb_size: int = 768):
        super(PatchEmbedding, self).__init__()
        self.patch_size = patch_size
        # 2D 이미지를 1D 벡터로 변환하는 Conv2D 레이어를 정의합니다.
        self.proj = nn.Conv2d(in_channels, emb_size, kernel_size=patch_size, stride=patch_size)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.proj(x)  # 이미지를 패치로 나누고 각 패치를 벡터로 변환합니다.
        x = x.flatten(2)  # 2D 형태의 패치들을 1D로 평탄화합니다.
        x = x.transpose(1, 2)  # 차원의 순서를 변경하여 [배치 크기, 패치 수, 임베딩 크기] 형태로 만듭니다.
        return x

• 이미지를 여러 개의 작은 사각형(패치)로 잘라내어 각 패치를 벡터(일렬로 늘어선 숫자들의 리스트)로 바꾸는 모듈입니다.
• init
  • 이 메서드는 PatchEmbedding 클래스의 인스턴스가 생성될 때 호출되는 초기화 함수입니다. 여기서는 필요한 변수들을 설정하고, 이미지를 패치로 분할하고 해당 패치를 임베딩 벡터로 변환하는 역할을 하는 convolutional layer를 정의합니다.
    • in_channels: 입력 이미지의 채널 수입니다. 일반적인 RGB 이미지의 경우, 이 값은 3이 됩니다 (Red, Green, Blue).
    • patch_size: 이미지를 분할할 때의 각 패치의 크기입니다. 예를 들어, patch_size가 16이라면, 이미지는 16x16 크기의 패치로 분할됩니다.
    • emb_size: 각 패치를 변환한 후의 임베딩 벡터의 크기입니다.
  • nn.Conv2d는 여기서 특별한 용도로 사용됩니다. 보통 컨볼루션 레이어는 이미지의 지역적 특징을 추출하는 데 사용되지만, 이 경우에는 이미지를 패치로 분할하고, 해당 패치를 임베딩 벡터로 변환하는 역할을 합니다. kernel_size와 stride를 patch_size로 설정하면, 이 컨볼루션 레이어는 non-overlapping 패치를 생성하게 됩니다.
  • 즉, PatchEmbedding 클래스의 초기화 함수는 이미지를 고정된 크기의 패치로 분할하고, 각 패치를 임베딩 벡터로 변환할 준비를 합니다.

Q. emb_size 가 768인 이유?

1. BERT와의 연관성: 768은 BERT라는 유명한 텍스트 트랜스포머 모델의 base 버전에서의 임베딩 크기와 같습니다. BERT의 이러한 설정은 좋은 성능을 보였기 때문에, 다른 트랜스포머 기반 모델들에서도 비슷한 설정을 사용하는 경향이 있습니다.
2. 계산 효율성: 임베딩 크기는 모델의 전체 파라미터 수와 연산량에 큰 영향을 미칩니다. 768은 대규모 모델에서도 계산상 효율적이면서 좋은 성능을 보장할 수 있는 값으로 알려져 있습니다.
3. Empirical Results: 구조나 하이퍼파라미터를 결정할 때, 연구자들은 다양한 설정을 실험으로 테스트합니다. 768과 같은 값을 선택하는 것은 그 값에서 좋은 성능을 얻었다는 경험적 결과에 기반한 것일 수 있습니다.
4. 유연성: 768은 2의 거듭제곱 값인 512와 1024 사이에 위치합니다. 이를 통해 연산의 효율성과 모델의 표현력 사이에 균형을 맞출 수 있습니다.

그렇지만, emb_size의 값을 768로 설정하는 것은 필수적인 것은 아닙니다. 다른 프로젝트나 응용 프로그램에서는 다른 크기의 임베딩을 사용할 수 있으며, 최적의 값을 찾기 위해 실험을 통해 다양한 설정을 시도해 볼 수 있습니다.

Q. x.transpose(1, 2)

x.transpose(1,2) 하면 1행 2열인데 어떻게 [배치크기, 패치수, 임베딩 크기] 3가지가 들어갈 수 있는지 의문이였다.

→ x.transpose(1, 2)에서 transpose 메서드의 인자 1과 2는 텐서의 차원을 가리킵니다. 이 말은 1번째 차원과 2번째 차원을 바꾼다는 의미입니다, 행과 열의 개념과는 조금 다르게 생각해야 합니다. 텐서에서 차원의 인덱스는 0부터 시작합니다.

예를 들어, 텐서 x의 형태(shape)가 [배치 크기, 임베딩 크기, 패치 수]라고 가정하면, x.transpose(1, 2)를 호출하면 차원 1 (임베딩 크기)와 차원 2 (패치 수)가 바뀌게 됩니다. 결과적으로 텐서의 형태는 [배치 크기, 패치 수, 임베딩 크기]가 됩니다.

따라서, x.transpose(1, 2)는 [배치 크기, 임베딩 크기, 패치 수] 형태의 텐서를 [배치 크기, 패치 수, 임베딩 크기]로 바꾸는 것입니다.

간단히 말해서, transpose는 지정된 두 차원을 교환합니다.

Q. 차원을 바꿔주는 이유?

x.transpose(1, 2)를 사용하여 [배치 크기, 임베딩 크기, 패치 수] 형태의 텐서를 [배치 크기, 패치 수, 임베딩 크기]로 바꾸는 것은 주로 텐서의 형태를 특정 연산이나 모델의 다른 부분에 맞춰주기 위한 것입니다.

Vision Transformer (ViT)에서는 각 이미지 패치의 임베딩을 시퀀스의 요소로 취급합니다. 따라서 텐서의 형태를 [배치 크기, 패치 수, 임베딩 크기]로 변경하는 것은 이후의 트랜스포머 인코더와 호환되도록 시퀀스의 형태를 준비하는 과정입니다.

구체적으로:

1. 시퀀스 처리: 트랜스포머 구조는 일반적으로 시퀀스 데이터를 처리합니다. 여기서 '시퀀스'는 이미지 내의 패치들을 의미합니다. 따라서, 패치의 수가 시퀀스의 길이와 같게 되고, 각 시퀀스 요소(즉, 각 패치)는 해당 임베딩 크기의 벡터로 표현됩니다.
2. 트랜스포머 호환성: 트랜스포머 인코더는 시퀀스의 각 요소에 대해 독립적으로 작동합니다. 입력 텐서의 형태가 [배치 크기, 시퀀스 길이, 피쳐 차원]이어야 합니다. 여기서 '시퀀스 길이'는 패치의 수와 같고, '피쳐 차원'은 임베딩 크기와 같습니다.
3. 포지셔널 임베딩: 트랜스포머는 시퀀스의 순서 정보를 자동으로 인식하지 않습니다. 따라서 포지셔널 임베딩을 통해 위치 정보를 추가합니다. 텐서의 형태를 [배치 크기, 패치 수, 임베딩 크기]로 변경하면, 포지셔널 임베딩을 각 패치의 임베딩에 간편하게 더할 수 있습니다.

요약하면, 텐서의 형태를 변경하는 것은 이후의 트랜스포머 인코더와 같은 모델의 구성 요소와의 호환성을 위한 것입니다.

---

3. 트랜스포머 인코더 블록

class TransformerEncoder(nn.Module):
    def __init__(self, emb_size: int = 768, drop_p: float = 0.1, forward_expansion: int = 4, num_heads: int = 8):
        super(TransformerEncoder, self).__init__()
        self.norm1 = nn.LayerNorm(emb_size)  # 임베딩 벡터를 정규화하는 레이어
        self.attention = nn.MultiheadAttention(embed_dim=emb_size, num_heads=num_heads)  # 멀티 헤드 어텐션 메커니즘
        self.norm2 = nn.LayerNorm(emb_size)  # 임베딩 벡터를 정규화하는 레이어
        # 두 개의 선형 레이어로 구성된 MLP (피드포워드 네트워크)
        self.mlp = nn.Sequential(
            nn.Linear(emb_size, forward_expansion * emb_size),
            nn.GELU(),
            nn.Linear(forward_expansion * emb_size, emb_size),
        )
        self.drop = nn.Dropout(drop_p)  # 과적합 방지를 위한 드롭아웃 레이어

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Attention block
        attn_output, _ = self.attention(x, x, x)  # 어텐션 메커니즘을 사용하여 정보를 결합
        x = x + self.drop(attn_output)  # 원래의 입력과 어텐션 출력을 더함
        x = self.norm1(x)  # 결과를 정규화
        
        # MLP block
        mlp_output = self.mlp(x)  # MLP를 통해 정보를 변환
        x = x + self.drop(mlp_output)  # 원래의 입력과 MLP 출력을 더함
        x = self.norm2(x)  # 결과를 정규화
    
        return x

• 트랜스포머는 정보를 처리하는 방식을 묘사하는 모델입니다. 여기서는 트랜스포머의 한 부분인 인코더를 정의하였습니다.
• 이 인코더는 주어진 정보 간의 관계를 학습합니다.
• init
  • emb_size: 임베딩 벡터의 차원 수입니다. 이전의 VisionTransformer 및 PatchEmbedding 설명에서 본 것처럼, 각 이미지 패치는 이 크기의 벡터로 임베딩됩니다.
  • drop_p: 드롭아웃 비율입니다. 드롭아웃은 과적합을 방지하는 데 사용되는 규제 방법입니다. drop_p는 제외될 노드의 확률을 나타냅니다.
  • forward_expansion: 피드포워드 네트워크 내에서 임베딩 차원을 확장하는 데 사용되는 계수입니다. 예를 들어, emb_size가 768이고 forward_expansion이 4라면, 피드포워드 네트워크의 내부 차원은 768 * 4 = 3072가 됩니다.
  • num_heads: 멀티 헤드 어텐션 메커니즘에서 사용되는 '헤드'의 수입니다. 멀티 헤드 어텐션은 입력을 여러 부분으로 분할하여 동시에 처리하고, 이를 다시 합쳐서 더 다양한 정보를 포착하는 데 도움을 줍니다.
  TransformerEncoder의 __init__ 메서드는 이러한 매개변수를 사용하여 해당 인코더 블록의 구성요소를 초기화합니다. 코드 스니펫에 따른 나머지 부분도 중요한데, 이 부분은 멀티헤드 어텐션, 정규화, 피드포워드 네트워크 등의 주요 트랜스포머 구성요소를 초기화하는 코드를 포함하고 있을 것입니다.

Q. 임베딩 벡터를 정규화 하는 이유?

임베딩 벡터 정규화 (Layer Normalization):

Layer normalization은 주어진 데이터 내에서 각 피쳐의 통계적 분포를 정규화합니다. 즉, 각 임베딩 벡터 내에서 동일한 정규화를 적용합니다.

정규화의 주요 이유는 아래와 같습니다:

• 학습 안정성 향상: 정규화는 각 레이어의 입력 분포를 일정하게 유지하여 학습 과정에서의 발산을 줄이고, 학습 속도를 높입니다.
• 그래디언트 소실/폭발 문제 완화: 정규화는 역전파되는 그래디언트의 크기를 안정화시켜 그래디언트 소실 또는 폭발 문제를 완화합니다.

4. 전체 Vision Transformer 모델

class VisionTransformer(nn.Module):
    def __init__(self, in_channels: int = 3, patch_size: int = 16, emb_size: int = 768, img_size: int = 224,
                 num_blocks: int = 12, num_classes: int = 1000, drop_p: float = 0.1):
        super(VisionTransformer, self).__init__()
        
        # 이미지를 패치로 잘라내고, 각 패치를 벡터로 바꾸는 부분
        self.patch_embed = PatchEmbedding(in_channels, patch_size, emb_size)
        
        # 이미지가 패치로 얼마나 잘라지는지 계산
        num_patches = (img_size // patch_size) ** 2
        
        # 분류를 돕는 특별한 토큰 (CLS Token)
        self.cls_token = nn.Parameter(torch.randn(1, 1, emb_size))
        
        # 각 패치의 위치 정보를 담은 임베딩
        self.pos_embed = nn.Parameter(torch.randn(1, num_patches + 1, emb_size))
        
        # 트랜스포머 인코더 블록을 여러 개 만듭니다.
        self.blocks = nn.ModuleList([
            TransformerEncoder(emb_size, drop_p)
            for _ in range(num_blocks)
        ])
        
        # 최종 출력을 위한 레이어
        self.mlp_head = nn.Sequential(
            nn.LayerNorm(emb_size),  # 벡터 정규화
            nn.Linear(emb_size, num_classes)  # 분류를 위한 선형 레이어
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B = x.shape[0]  # 입력 이미지의 배치 크기

        # 이미지를 패치로 나누고 벡터로 변환
        x = self.patch_embed(x)
        
        # 각 이미지에 CLS 토큰 추가
        cls_tokens = self.cls_token.expand(B, -1, -1)
        x = torch.cat((cls_tokens, x), dim=1)
        
        # 위치 정보 추가
        x += self.pos_embed
        
        # 모든 트랜스포머 블록을 순차적으로 거치게 합니다.
        for block in self.blocks:
            x = block(x)
            
        # CLS 토큰의 정보만 가져옴 (이 정보를 사용해 이미지를 분류)
        x = x[:, 0]
        
        # 이미지 분류를 위해 마지막 레이어를 거침
        x = self.mlp_head(x)
        
        return x  # 분류된 결과 반환

• 앞서 정의한 패치 임베딩과 트랜스포머 인코더를 조합하여 전체 이미지를 처리하고 분류하는 작업을 수행하는 모델입니다.
• cls_token: 분류를 위한 특별한 토큰입니다. 이것은 트랜스포머가 어떤 카테고리에 이미지가 속하는지 판단하는 데 도움을 줍니다.
• pos_embed: 각 패치의 위치 정보를 제공하는 임베딩입니다. 위치 정보는 트랜스포머가 이미지의 구조를 이해하는 데 중요합니다.
• blocks: 트랜스포머 인코더를 여러 개 쌓아 올린 것입니다. 이로 인해 모델은 이미지의 복잡한 패턴을 학습할 수 있습니다.
• mlp_head: 최종적으로 이미지가 어떤 카테고리에 속하는지 결정하는 부분입니다.
• init
  • in_channels: 입력 이미지의 채널 수입니다. 대부분의 경우 RGB 이미지를 다루기 때문에 기본값은 3입니다.
  • patch_size: 입력 이미지를 분할할 때 사용되는 패치의 크기입니다. 이미지는 patch_size x patch_size 크기의 패치들로 분할됩니다.
  • emb_size: 각 패치가 변환될 임베딩의 크기입니다.
  • img_size: 입력 이미지의 크기입니다. Vision Transformer는 일반적으로 정사각형 이미지를 다루기 때문에 이 값은 이미지의 높이와 너비 모두를 의미합니다.
  • num_blocks: Transformer 인코더 블록의 수입니다. 각 블록은 멀티헤드 어텐션과 피드포워드 네트워크를 포함합니다.
  • num_classes: 최종 분류될 클래스의 수입니다. 예를 들면, ImageNet 데이터셋에 대한 분류 작업의 경우 이 값은 1000입니다.
  • drop_p: 드롭아웃 비율을 지정하는 값입니다.
  이 메서드의 내용 부분 (...에 해당하는 부분)은 Vision Transformer의 핵심 구성 요소를 초기화하는 코드를 포함합니다:
  1. Patch Embedding: 이미지를 여러 패치로 분할하고 각 패치를 벡터로 임베딩하는 역할을 합니다.
  2. Class Token: 모든 입력 시퀀스의 시작 부분에 추가되는 특별한 토큰입니다. 최종 분류 작업을 위해 주로 사용됩니다.
  3. Positional Embedding: 시퀀스 내의 각 패치의 위치 정보를 추가하는 역할을 합니다.
  4. Transformer Encoder Blocks: 주어진 수(num_blocks)만큼의 트랜스포머 인코더 블록들의 시퀀스입니다.
  5. MLP Head: 최종 분류를 위한 멀티레이어 퍼셉트론입니다.
  이 __init__ 메서드는 이러한 구성 요소들을 초기화하고, 주어진 매개변수와 설정에 따라 Vision Transformer 모델을 구성합니다.