1D vector로 만들어진 patch들을 Linear Projection을 통해 768 차원의 각 patch의 embedding vector로 표현.
class PatchEmbedding(nn.Module):
def __init__(self, in_channels: int = 3, patch_size: int = 16, emb_size: int = 768):
self.patch_size = patch_size
super().__init__()
'''
self.projection = nn.Sequential(
# break-down the image in s1 x s2 patches and flat them
Rearrange('b c (h s1) (w s2) -> b (h w) (s1 s2 c)', s1=patch_size, s2=patch_size),
nn.Linear(patch_size * patch_size * in_channels, emb_size)
)
'''
self.projection = nn.Sequential(
# using a conv layer instead of a linear one -> performance gains
nn.Conv2d(in_channels, emb_size, kernel_size=patch_size, stride=patch_size),
Rearrange('b e (h) (w) -> b (h w) e'),
)
def forward(self, x: Tensor) -> Tensor:
x = self.projection(x)
return x
각 patch embedding에 classification token, position embedding 추가.
Transformer Encoder
Transformer Output
classification
■학습을 통해 결정되는 Parameter
Classification token
classification을 위해 사용되는 token
class PatchEmbedding(nn.Module):
def __init__(self, in_channels: int = 3, patch_size: int = 16, emb_size: int = 768):
self.patch_size = patch_size
super().__init__()
self.projection = nn.Sequential(
# using a conv layer instead of a linear one -> performance gains
nn.Conv2d(in_channels, emb_size, kernel_size=patch_size, stride=patch_size),
Rearrange('b e (h) (w) -> b (h w) e'),
)
self.cls_token = nn.Parameter(torch.randn(1,1, emb_size))
def forward(self, x: Tensor) -> Tensor:
b, _, _, _ = x.shape
x = self.projection(x)
cls_tokens = repeat(self.cls_token, '() n e -> b n e', b=b)
# prepend the cls token to the input
x = torch.cat([cls_tokens, x], dim=1)
return x
Position embedding
patch의 위치 정보를 가지고 있는 embedding
position embedding을 사용 안하면?
이미지 입력 → 일전 크기의 patch로 자름 → patch를 sequence로 생각하여 transformer Encoder로 입력 → 각 patch가 어떤 위치에서 왔는지 위치 정보가 손실됨. → 학습 안됨
(classification token, patch embedding) + position embedding 이 Transformer Encoder로 입력됨.
classification token, position token은 학습에 의해 결정되는 parameter입니다.
class PatchEmbedding(nn.Module):
def __init__(self, in_channels: int = 3, patch_size: int = 16, emb_size: int = 768, img_size: int = 224):
self.patch_size = patch_size
super().__init__()
self.projection = nn.Sequential(
# using a conv layer instead of a linear one -> performance gains
nn.Conv2d(in_channels, emb_size, kernel_size=patch_size, stride=patch_size),
Rearrange('b e (h) (w) -> b (h w) e'),
)
self.cls_token = nn.Parameter(torch.randn(1,1, emb_size))
self.positions = nn.Parameter(torch.randn((img_size // patch_size) **2 + 1, emb_size))
def forward(self, x: Tensor) -> Tensor:
b, _, _, _ = x.shape
x = self.projection(x)
cls_tokens = repeat(self.cls_token, '() n e -> b n e', b=b)
# prepend the cls token to the input
x = torch.cat([cls_tokens, x], dim=1)
# add position embedding
x += self.positions
return x
시각화
cos = torch.nn.CosineSimilarity(dim=1, eps=1e-6)
fig = plt.figure(figsize=(10,10))
fig.suptitle("Visualization of position embedding similarities", fontsize=24)
for i in range(1, pos_embed.shape[1]):
sim = F.cosine_similarity(pos_embed[0, i:i+1], pos_embed[0, 1:], dim=1)
sim = sim.reshape((14, 14)).detach().cpu().numpy()
ax = fig.add_subplot(14, 14, i)
ax.axes.get_xaxis().set_visible(False)
ax.axes.get_yaxis().set_visible(False)
ax.imshow(sim)
■Architecture
Transformer Encoder
ViT는 Vanilla Transformer(Attent is all you need)구조와는 다른 구조로 구성되어 있다.
Transformer는 layer를 깊게 쌓게되면 학습이 어렵다는 단점이 있다.
이런 단점을 극복하여 학습이 되게 하기 위해서는 layer normalization의 위치가 중요하다는 것이 후속 연구들을 통해 증명이 되었다.
ViT는 NLP 에서 보여진 layer를 깊게 쌓으면 학습이 어렵다는 문제와 이를 극복하기 위해 layer normalization 위치가 중요하다는 연구를 받아드린 Vision Transformer이다.
기존 Transformer: Multi-head attention → normalization
weight metrix $W_Q$, $W_K$, $W_V$에 곱해져서 Query, Key, VAlue 값들이 계산됨(⭐️ weight metrix의 학습을 통해 attention이 학습)
Query, Key 값들의 dot product[torch.einsum('bhqd, bhkd -> bhqk', queries, keys)]를 통해 Similarity 계산
Softmax를 취해 0 ~ 1의 attention score계산
계산된 attention score에 Value를 곱해 최종적인 output을 얻음.
Multi-Head Self-Attention
ViT base model에서는 self-attention 12번 수행.
768차원 → Multi-head self-attention 통과 → 64차원으로 감소 → self-attention을 12번 수행하니 768차원으로 증가 → 결국 차원수가 동일하게 유지됨
class MultiHeadAttention(nn.Module):
def __init__(self, emb_size: int = 768, num_heads: int = 8, dropout: float = 0):
super().__init__()
self.emb_size = emb_size
self.num_heads = num_heads
self.keys = nn.Linear(emb_size, emb_size)
self.queries = nn.Linear(emb_size, emb_size)
self.values = nn.Linear(emb_size, emb_size)
self.att_drop = nn.Dropout(dropout)
self.projection = nn.Linear(emb_size, emb_size)
self.scaling = (self.emb_size // num_heads) ** -0.5
def forward(self, x : Tensor, mask: Tensor = None) -> Tensor:
# split keys, queries and values in num_heads
queries = rearrange(self.queries(x), "b n (h d) -> b h n d", h=self.num_heads)
keys = rearrange(self.keys(x), "b n (h d) -> b h n d", h=self.num_heads)
values = rearrange(self.values(x), "b n (h d) -> b h n d", h=self.num_heads)
# sum up over the last axis
energy = torch.einsum('bhqd, bhkd -> bhqk', queries, keys) # batch, num_heads, query_len, key_len
if mask is not None:
fill_value = torch.finfo(torch.float32).min
energy.mask_fill(~mask, fill_value)
att = F.softmax(energy, dim=-1) * self.scaling
att = self.att_drop(att)
# sum up over the third axis
out = torch.einsum('bhal, bhlv -> bhav ', att, values)
out = rearrange(out, "b h n d -> b n (h d)")
out = self.projection(out)
return out
