수달이네 기술 블로그

1. ViT(Vision Transformer 본문

AI공부/Vision

1. ViT(Vision Transformer

슬픈 수달이 2026. 3. 9. 13:03

ViT(Vision Transformer)

이미지 인식 문제를 기존의 CNN이 아닌 Transformer구조로 해결한 모델

  • 이미지를 작은 Patch로 나누어 이를 토큰으로 취급
  • 해당 토큰을 Self-Attention으로 학습하여 문맥을 이해한다.

장점

  • 멀리 떨어진 영역 간의 상관관계를 한번에 파악 가능
  • 데이터가 충분하다면 CNN기반 모델을 능가할 수 있음.
  • 이미지 크기가 달라져도 패치와 그 수를 변경하는 등으로 유연하게 대응가능
  • 이미지를 “시퀀스”로써 바라본다는 관점 전환으로 비전, 자연어 간의 경계를 허물었다.

구조

  1. 이미지 입력(input image)
    모델에 원본 이미지를 입력. CNN과 달리 convolution연산을 진행하지 않고 그대로 사용.
  2. 이미지를 패치단위로 분할(Divide Image into Pathces)
    기존 이미지(2차원 행렬) > 분할 후(패치 집합)
    이미지를 고정 크기의 패치로 분할 후 독립적으로 처리
  3. 패치단위 토큰으로 전환 (Convert Patches into Tokens)
    이미지 패치를 flatten 연산을 거치고, 선형 변환하여 고정 차원의 벡터로 변환된다.(임베딩 벡터)
    해당 토큰을 자연어에서의 토큰과 같이 처리함.
  4. 인코더(Transformer Encoder)
    토큰이 Transformer Encoder에 입력된 후 인코더에서 연산을 진행
    • Self-Attention연산 > 패치 내 토큰 간의 관계(문맥)를 학습(Learning Global Context)
      • CNN은 국소 영역에 의존함
    • Feed-Forward Network로 비선형으로 변환
    • 잔차연결과 Layer Normalization으로 학습 안정성, 정보 보존
    인코더에 들어가는 연산은 모두 병렬처리 된다.
  5. 출력(분류 결과)
    분류 헤드(Classification head)를 거쳐 최종 예측 결과를 생성함.
    • 패치간의 정보를 종합해 판단

구현

import numpy as np

import torch
import torch.nn as nn
import torch.nn.functional as F

이미지 생성

image = np.random.rand(8, 8, 3).astype(np.float32)  # 생성
image = image.reshape(2, 4, 2, 4, 3) # 분할
image = image.transpose(0, 2, 1, 3, 4) # 정렬
patches = image.reshape(-1, 4, 4, 3) # 패치 설정

print("patches:", patches.shape)
  • 8x8 이미지를 랜덤 생성
  • 이후 2, 4, 2, 4, 3으로 재배치 > 가로 세로 2패치씩 나눠줌.
  • 나눈 패치를 축을 바꿔서 정렬해줌

하이퍼 파라미터 설정

patches = torch.tensor(patches) # 패치를 텐서로 변환

embedding_dim = 32 # 벡터 차원 수
num_heads = 4    # 헤드 개수
num_transformer_layers = 2 # 어텐션 + (MLP)다층 퍼셉트론층수
num_classes = 10 # 분할 클래스 개수
mlp_dim = 256 # MLP차원수
  • 패치 > 텐서로 변환(신경망의 입력으로 사용)

임베딩 차원에 매핑

def patch_embedding(patches, embedding_dim):
    N, Ph, Pw, C = patches.shape
    patch_dim = Ph * Pw * C 
    patch_flat = patches.reshape(N, patch_dim)
    patch_flat = patch_flat.unsqueeze(0)

    proj = nn.Linear(patch_dim, embedding_dim)
    tokens = proj(patch_flat)

    return tokens, proj

tokens, patch_proj = patch_embedding(patches, embedding_dim)
print("tokens:", tokens.shape) 

#tokens: torch.Size([1, 4, 32])
  • 패치를 벡터로 변환

ViT 실행

B, N, D = tokens.shape # B:배치크기, N: 패치개수, D:임베딩 차원

cls_token = nn.Parameter(torch.zeros(1, 1, D))  # 분류를 위한 토큰
x = torch.cat([cls_token.expand(B, -1, -1), tokens], dim=1)

pos_embed = nn.Parameter(torch.randn(1, x.size(1), D) * 0.02)  # 포지셔널 임베딩(간략화)
x = x + pos_embed 

encoder_layer = nn.TransformerEncoderLayer( # 인코더
    d_model=D,
    nhead=num_heads,
    dim_feedforward=mlp_dim,
    batch_first=True
)
transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_transformer_layers)

z = transformer_encoder(x) # 입력토큰 x를 인코더에 넣음

classification_layer = nn.Linear(D, num_classes)
logits = classification_layer(z[:, 0, :]) # 해당 결과로 분류
output = F.log_softmax(logits, dim=-1)
pred_class = output.topk(1, dim=-1).indices.item() # 클래스 확률 계산
print("class :", pred_class) # 출력
  1. 패치를 임베딩 한후
  2. cls토큰을 추가하고, positional embedding한 후
  3. 인코더를 통과한 후
  4. 분류하는 방식

위처럼 ViT의 작동을 간략화해 보았다.

'AI공부 > Vision' 카테고리의 다른 글

3. AutoEncoder2(Sparse AE, Denoising AE, AutoEncoder + CIFAR10)  (0) 2026.03.11
2. AutoEncoder(오토인코더)  (0) 2026.03.10