CV와 NLP에서 공통적으로 나타나는 문제 중 하나는 Long-range dependency 문제이다. CNN의 convolution filter는 local 정보에 집중한다. 이는 작은 receptive field 내에서만 정보를 처리하고 깊은 층을 쌓을수록 정보를 통합하는 방식이다. 이는 필터 크기를 늘린다고 해도, 근본적인 global 정보를 처리하는데 한계가 있다.
-> CNN 대신 Transformer 구조를 사용할 수 없을까?
image의 픽셀을 token으로 사용하자. -> (256, 256) -> 256*256 = 65536 = 토큰 개수 -> 막대한 양의 연산이 필요하다. 따라서 모든 픽셀에 대해 attention weight를 계산하는 것은 매우 비효율적이다.
-> N개의 픽셀을 묶어서 token으로 사용하자 -> N * N patch를 만들어서 token으로 사용하자
Vision Transformer 흐름
입력: 이미지를 작은 패치들로 나누고, 각 패치를 임베딩 벡터로 변환한 후, 패치의 순서를 알리기 위해 Positional Encoding을 더한다.
Encoder: 변환된 패치들을 입력으로 받아, Self-Attention을 수행하여 패치 간의 관계를 학습시키며, Encoder block을 여러 번 반복하여 각 패치의 전역적 의미를 이해.
[CLS]토큰: 이미지 전체의 정보를 요약하기 위해 [CLS]토큰을 사용 -> 이는 최종적으로 이미지 분류에 사용.
classification: [CLS] 토큰에 담긴 정보를 이용해 최종 분류 결과를 출력
Patch embedding
코드에서 224*224 이미지를 16*16패치로 나누고, embedding을 구현하기 위해 간단하게 Convolution layer를 이용하면 쉽게 구현할 수 있다.
처음에 입력 이미지 3개의 RGB채널로부터 embed_size인 768개의 커널로 filter, stride를 patch_size만큼 설정하면 각 patch에 커널이 찍히며 각 patch들이 768차원으로 매핑되므로, 각 패치의 특징이 잘 반영된다.
patch_embedding + Positional Encoding 값을 transformer와 다르게 Layer Norm을 통과한 값을 인풋으로 받는다. 인풋을 각각 projection 시켜 query, key, value를 만들고, head 수만큼 split을 진행한다. -> 앙상블 효과
query와 key를 통해 Attention score를 계산 후 Value를 곱함으로써 각각의 Key가 가진 정보를 통합
각 Query를 대표하는 Attention Value를 얻은 후, 각 Head의 Attention score를 종합하여 patch로 분할된 이미지를 patch별 잘 이해할 수 있게 된다.
class MultiHeadSelfAttention(nn.Module):
def __init__(self, model_dim=768, num_heads=8, attn_drop=0., out_drop=0.):
super().__init__()
assert model_dim % num_heads ==0
self.model_dim = model_dim
self.num_heads = num_heads
self.head_dim = model_dim // num_heads
# 선형 변환 Linear layer 정의
self.q_linear = nn.Linear(model_dim, model_dim)
self.k_linear = nn.Linear(model_dim, model_dim)
self.v_linear = nn.Linear(model_dim, model_dim)
self.output_linear = nn.Linear(model_dim, model_dim)
# Attention에 dropout 적용
self.attn_drop = nn.Dropout(attn_drop)
self.output_drop = nn.Dropout(out_drop)
def split_heads(self, x):
"""
(B, N, model_dim) -> (B, N, num_heads, head_dim)
transpose(1,2) -> (B, num_heads, N, head_dim)
"""
B, N, C = x.size()
return x.view(B, N, self.num_heads, self.head_dim).transpose(1,2)
def combine_heads(self, x):
"""
여러 헤드에서 계산된 결과를 다시 결합
(B, n_heads, N, head_dim) -> (B, N, model_dim)
"""
B, _, N, head_dim = x.size()
return x.transpose(1,2).contiguous().view(B, N, self.model_dim)
def forward(self, x):
# query, key, value를 선형 변환 후 헤드별로 분할 (split_heads로 n_heads 차원 추가)
q = self.split_heads(self.q_linear(x))
k = self.split_heads(self.k_linear(x))
v = self.split_heads(self.v_linear(x)) # (B, num_heads, N, head_dim)
# Attention score 게산
attn = torch.matmul(q,k.transpose(-2,-1)) * math.sqrt(self.head_dim)
attn = attn.softmax(dim=-1)
attn = self.attn_drop(attn) # (B, num_heads, N, N)
# Attention 확률과 Value의 행렬 곱을 계산
x = torch.matmul(attn, v) # (B, num_heads, N, head_dim)
x = self.output_linear(self.combine_heads(x)) # (B, N, model_dim)
x = self.output_drop(x)
return x
Transformer Encoder Block
위에서 구현한 Multi-Head Attention 를 통해 얻은 결과를 Feedforward Network에 통과하는 것을 1개의 Block라고 한다.
Drop Path : 쉽게 말하면 Dropout은 각 뉴련 수준에서 무작위로 연결을 끊지만, Drop Path는 특정 경로 전체를 끊는다. -> 일반화 성능을 높인다.
Drop Path 구현 방식: 인풋 차원만큼 (batch, 1, 1)로 shape을 만들고 bernoulli 분포를 이용하여 설정된 확률만큼 0과 1로 이루어진 텐서를 생성하여 경로 전체를 무작위로 drop을 한다. 이때 x.div(keep_prob)로 스케일링하는 이유는 Train시에만 작동하므로, Test시 Train 단계보다 출력값이 커지기 때문에 이를 맞춰주기 위해서이다.
class EncoderBlock(nn.Module):
"""Transformer Encoder block"""
def __init__(self, model_dim, num_heads, mlp_ratio=4, mlp_drop = 0., attn_drop=0, drop_path=0.):
super().__init__()
self.norm1 = nn.LayerNorm(model_dim)
self.attn_layer = MultiHeadSelfAttention(
model_dim, num_heads, attn_drop=attn_drop, out_drop = mlp_drop
)
self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
self.norm2 = nn.LayerNorm(model_dim)
mlp_hidden = model_dim * mlp_ratio
self.mlp = Mlp(
in_features=model_dim, hidden_features=mlp_hidden, act_layer=nn.GELU, drop=mlp_drop
)
def forward(self, x):
x = self.attn(self.norm1(x))
x = x + self.drop_path(x)
x = x + self.drop_path(self.mlp(self.norm2(x)))
return x
def drop_path(x, drop_prob: float = 0., training: bool = False):
""" Dropout과 비슷하게 overfitting을 방지 """
if drop_prob == 0. or not training:
return x
keep_prob = 1 - drop_prob
shape = (x.shape[0],) + (1,) * (x.ndim - 1) # (batch, 1, 1)
random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
output = x.div(keep_prob) * random_tensor
return output
class DropPath(nn.Module):
def __init__(self, drop_prob: float = 0.):
super().__init__()
self.drop_prob = drop_prob
def forward(self, x):
return drop_path(x, self.drop_prob, self.training)
class Mlp(nn.Module):
"""Transformer Encoder의 feedforward Network"""
def __init__(self, in_features, hidden_features, act_layer=nn.GELU, drop=0.):
super().__init__()
self.fc1 = nn.Linear(in_features, hidden_features)
self.act = act_layer()
self.fc2 = nn.Linear(hidden_features, in_features)
self.drop = nn.Dropout(drop)
def forward(self, x):
x = self.act(self.fc1(x))
x = self.drop(x)
x = self.drop(self.fc2(x))
return x
VisitionTransformer
이 class는 지금까지 구현한 VisitionTransformer를 구현하기 위해 조각들을 다 합쳐서 visitionTransformer를 정의하는 것이다.
전반적인 흐름은 들어온 input으로 기본값 3x224x224를 받아 앞서 정의한 patch_embedding을 거쳐 positional_encoding을 더하고, Transformer Encoder Block를 depth=12만큼 통과하는 것이다. 여기서 주요깊에 봐야 할 것은 interpolate_pos_encoding이다.
interpolate_pos_encoding: 일반적으로는 nn.Parameter로 patch_embedding shape으로 만들어주면 되지만, patch_size=16으로 224x224인 h,w로 Pretrain 된 Positional encoding을 사용할 때, 우리의 input이 224x224가 아닌, 예를 들어 320x320인 크기를 patch_size=16으로 자를 때, 20x20의 Positional encoding이 나오기 때문에, 미리 학습된 16x16의 Postitional_encoding을 사용하기 위해, interpolation을 진행하여 20x20으로 만들어준다.
interpolate_pos_encoding코드 중, w0, h0 = w0 + 0.1, h0 + 0.1 을 한 이유는 nn.functional.interpolate()의 scale_factor에서 정수로 나눠 떨어지지 않을 때, 만약 값이 19.9999이면 우리가 원하는 값은 정수 20인데, int()로 변환되어 19가 반환되기 때문에, 0.1을 더해줘서 20.9999로 int()를 했을 때, 우리가 원하는 20을 얻을 수 있기 때문이다.
class VisionTransformer(nn.Module):
""" Vision Transformer 정의 """
def __init__(
self,img_size=[224], patch_size=16, in_chans=3, num_classes=0, embed_dim=768,
depth=12, num_heads=12, mlp_ratio=4., drop_rate=0., attn_drop_rate=0., drop_path_rate=0.,
norm_layer=nn.LayerNorm,**kwargs
):
super().__init__()
# Embedding size
self.embed_dim = embed_dim
# From image to patch
self.patch_embed = PatchEmbed(img_size=img_size[0], patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
# 인풋 image와 patch 크기로부터 계산된 patch 갯수 (224/16 = 14 -> 가로, 세로 14개씩 patch로 총 196개의 patch가 생성)
num_patches = self.patch_embed.num_patches
# CLS Token 생성,
# batch_size, num_patch를 1 로 설정 후 input에 맞춰 브로드캐스트
self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
# Positional embedding 생성, 학습가능한 파라미터 설정 (각 patch + class token)
self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
# Postional embedidng 결과에 dropout을 적용
self.pos_drop = nn.Dropout(p=drop_rate)
# Drop_path의 Drop확률을 depth에 따라 차등적용
dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]
# Transformer의 encoder block
self.blocks = nn.ModuleList([
Block(
dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio,
drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer)
for i in range(depth)])
self.norm = norm_layer(embed_dim)
# Class token과 positional embedding의 parameter 초기화
trunc_normal_(self.pos_embed, std=.02)
trunc_normal_(self.cls_token, std=.02)
# Layernorm과 classification head의 parameter 초기화
self.apply(self._init_weights)
def _init_weights(self, m):
"""Layernorm과 classification head의 parameter 초기화하는 함수"""
if isinstance(m, nn.Linear):
trunc_normal_(m.weight, std=.02)
if isinstance(m, nn.Linear) and m.bias is not None:
nn.init.constant_(m.bias, 0)
elif isinstance(m, nn.LayerNorm):
nn.init.constant_(m.bias, 0)
nn.init.constant_(m.weight, 1.0)
def interpolate_pos_encoding(self, x, w, h):
"""Pre-trained된 positional encoding을 interpolation을 사용하여 고해상도 이미지에 적용하기 위한 함수"""
npatch = x.shape[1] - 1 # 197 - 1 = 196
N = self.pos_embed.shape[1] - 1 # 196
# Interpolation 할 필요가 없으면 interpoltation 없이 postional embedding 적용
if npatch == N and w == h:
return self.pos_embed # (1, 197, 768)
# Interpolation을 위한 코드
class_pos_embed = self.pos_embed[:, 0] # CLS token해당 부분 따로 추출 - (1,768)
patch_pos_embed = self.pos_embed[:, 1:] # CLS 제외한 Patch에 대한 pos_enc 추출(실제 이미지 부분) - (1, 196, 768)
dim = x.shape[-1] # dim = 768
# 입력 이미지가 320x320 일 때, 패치 크기는 16일 때
w0 = w // self.patch_embed.patch_size # 320 // 16 = 20
h0 = h // self.patch_embed.patch_size # 320 // 16 = 20
w0, h0 = w0 + 0.1, h0 + 0.1 # 20.1, 20.1 # # Interpolation 과정에서 rounding을 피하기 위한 작은 보정 ---- 만약 19.9999이면 우리가 원하는 값은 20인데 int()로 정수로 맞춰지기때문에 0.1을 더해서 우리가 원하는 20.999여서 20으로 맞추기 위해서 이다.
patch_pos_embed = nn.functional.interpolate(
patch_pos_embed.reshape(1, int(math.sqrt(N)), int(math.sqrt(N)), dim).permute(0, 3, 1, 2),
scale_factor=(w0 / math.sqrt(N), h0 / math.sqrt(N)), # => 14x14 에서 20x20으로 키우겠다.
mode='bicubic',
) # shape (1, 196, 768)->resahpe(1, 768, 14, 14) 보간-> (1, 768, 20, 20)
assert int(w0) == patch_pos_embed.shape[-2] and int(h0) == patch_pos_embed.shape[-1]
patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim) # (1, 400, 768)
return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1) # (1, 401, 768)
def prepare_tokens(self, x, interpolate_pos_encoding=False):
"""이미지 embedding후 classification token과 positional embedding을 추가로 적용하는 함수"""
# (batch, channel, width, height)
B, nc, w, h = x.shape # ex (32, 3, 224, 224)
# 이미지를 토큰화후 embedding
x = self.patch_embed(x) # (b, 196, 768) = (b, num_patch^2, embed_size)
# [CLS] 토큰을 batch size에 맞게 확장
cls_tokens = self.cls_token.expand(B, -1, -1) # (32, 1, 768)
# [CLS] 토큰을 embedding patch에 추가
x = torch.cat((cls_tokens, x), dim=1) # (32, 197, 768)
# 각 토큰에 positional_enocding 더하기
x = x + self.interpolate_pos_encoding(x, w, h) # (x, 224, 224)
# dropout 적용
x = self.pos_drop(x)
return x
def forward(self, x, return_all_patches=False):
# 이미지 -> patch화 -> embedding + positional_encoding
x = self.prepare_tokens(x)
# 정해진 depth만큼 encoder block 반복 통과
for blk in self.blocks:
x = blk(x)
# encoder output을 normalization
x = self.norm(x)
# 학습된 모든 patch가 필요한 경우 모든 patch의 latent feature map을 반환
if return_all_patches:
return x
# CLS Token의 latent feature map만 반환
else:
return x[:, 0]
Vit 모델 크기 조절
모델크기를 작게 하고싶으면, embed_dim, depth(block개수), num_heads, mlp_ratio와 같이 모델의 파라미터 수에 영향을 주는 값을 바꿔서 상황에 맞게 size조절하도록 한다.
def vit_tiny(patch_size=16, **kwargs):
model = VisionTransformer(
patch_size=patch_size, embed_dim=192, depth=12, num_heads=3, mlp_ratio=4,
qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
return model
def vit_small(patch_size=16, **kwargs):
model = VisionTransformer(
patch_size=patch_size, embed_dim=384, depth=12, num_heads=6, mlp_ratio=4,
qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
return model
def vit_base(patch_size=16, **kwargs):
model = VisionTransformer(
patch_size=patch_size, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4,
qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
return model
VisionTransformerWithLinear
분류 task라면, VisionTransformer에서 CLS Token의 latent feature map만 받아서 Linear()에 태워서 분류하고자 하는 class만큼 로짓을 계산을 한다.
class ViTWithLinear(nn.Module):
def __init__(self, base_vit, embed_dim=384, num_classes=10, **kwargs):
"""
Vision Transformer에 추가로 Linear Classifier를 붙인 class
Args:
base_vit (nn.Module): 사전 학습된 Vision Transformer 모델.
embed_dim (int): ViT에서 출력되는 임베딩의 차원. 기본값은 384.
num_classes (int): 분류할 클래스의 수. 기본값은 10.
"""
super().__init__()
# 사전 학습된 Vision Transformer 모델을 받음
self.base_vit = base_vit
# ViT의 출력 임베딩을 클래스 수에 맞게 변환하는 선형 계층 정의
self.fc = nn.Linear(embed_dim, num_classes)
def forward(self, x):
"""
Args:
x (torch.Tensor): 입력 이미지 텐서, 크기 (batch_size, channels, height, width).
Returns:
logits (torch.Tensor): 각 클래스에 대한 예측 점수.
"""
# 정의한 vit로 부터 image feature 추출
features = self.base_vit(x)
# feature 정규화
features = torch.nn.functional.normalize(features, dim=-1)
# classifier head를 거쳐 최종 logit 계산
logits = self.fc(features)
# 클래스별 예측 점수 반환
return logits
