[Vision Transformer] 코드 설명

컨텐트

Vision Transformer에 대한 코드 설명을 수행하고 싶습니다.
논문: An image is worth 16x16 words: Transformers for image recognition at scale
사용한 코드는 여기 있습니다.
https://github.com/lucidrains/vit-pytorch/blob/main/vit_pytorch/vit.py

Pre-Norm

class PreNorm(nn.Module):
    def __init__(self, dim, fn):
        super().__init__()
        self.norm = nn.LayerNorm(dim)
        self.fn = fn
    def forward(self, x, **kwargs):
        return self.fn(self.norm(x), **kwargs)

Transformer의 Sub-Layer에 사용되는 카테고리입니다.본가의 Transformer는 Post-정규, Vision Transformer는 Ple-정규fn 대입Multi-Head Attention과 Feed Forward Network을 사용하여 정규화된 값을 처리합니다.
※ 플레-정규를 채택한 이유: Learning Deep Transformer Models for Machine Translation

FeedForwardNetwork

class FeedForward(nn.Module):
    def __init__(self, dim, hidden_dim, dropout = 0.):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(dim, hidden_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, dim),
            nn.Dropout(dropout)
        )

    def forward(self, x):
        return self.net(x)

Transformer-Block의 2층에서 사용하는 반이다.특별히 설명할 것이 없다.

Attention 기구

class Attention(nn.Module):
    def __init__(self, dim, heads = 8, dim_head = 64, dropout = 0.):
        super().__init__()
        inner_dim = dim_head *  heads
        project_out = not (heads == 1 and dim_head == dim)

        self.heads = heads
        self.scale = dim_head ** -0.5

        self.attend = nn.Softmax(dim = -1)
        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)

        self.to_out = nn.Sequential(
            nn.Linear(inner_dim, dim),
            nn.Dropout(dropout)
        ) if project_out else nn.Identity()

    def forward(self, x):
        qkv = self.to_qkv(x).chunk(3, dim = -1)
        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), qkv)

        dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale

        attn = self.attend(dots)

        out = torch.matmul(attn, v)
        out = rearrange(out, 'b h n d -> b n (h d)')
        return self.to_out(out)

반입니다.다음 그림의 모델을 설치하고 있습니다.

참조: Attention is all you need

Attention

inner_dim = dim_head *  heads

__init__는 Query, Key, Value를 생성하는 라인 레이어의 출력 비트입니다.
멀티-Head Self-Attention을 사용했기 때문에 각 Q, K, V에 대해 헤드 수만 원하기 때문에 각 헤드의 차원수와 헤드 수(heads)를 곱한 값을 설정합니다.

self.scale = dim_head ** -0.5

\sqrt{D k}.

self.attend = nn.Softmax(dim = -1)

Attention을 구할 때 Softmax(QK^T/\sqrt{D k})가 사용하는 Softmax 함수입니다.

self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)

이 Linker 레이어는 [q, k, v]=zU{qkv} 의 U{qkv}에 부합합니다.출력에는 q, k, v 3개의 변수, inner가 필요합니다.dim*3이 됩니다.

self.to_out = nn.Sequential(
    nn.Linear(inner_dim, dim),
    nn.Dropout(dropout)
) if project_out else nn.Identity()

는 그림의 최종 층inner_dim에 부합된다.

Linear

qkv = self.to_qkv(x).chunk(3, dim = -1)

forward에서 수조를 3개로 나누어 3개 요소의 원조를 구성한다.

q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), qkv)

.chunk() 함수에서 매 헤드분할의 형식으로 배열을 변형시킨다.
구체적으로 (1651024)헤드 수는 16시(1,16,65,64)이다.

dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale
attn = self.attend(dots)

Softmax(QK^T/\sqrt{D k}) 계산이 진행 중입니다.

out = torch.matmul(attn, v)

SA(z)=Av 계산이 진행 중입니다.

out = rearrange(out, 'b h n d -> b n (h d)')
return self.to_out(out)

마지막으로 패치별로 변형하여 Linter 레이어로 가져옵니다.

Transformer Encoder

class Transformer(nn.Module):
    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout = 0.):
        super().__init__()
        self.layers = nn.ModuleList([])
        for _ in range(depth):
            self.layers.append(nn.ModuleList([
                PreNorm(dim, Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout)),
                PreNorm(dim, FeedForward(dim, mlp_dim, dropout = dropout))
            ]))
    def forward(self, x):
        for attn, ff in self.layers:
            x = attn(x) + x
            x = ff(x) + x
        return x

Transformer Enceoder입니다.위에서 만든 종류를 사용하여 간단하게 실현하였다.논문에서 이쪽으로.

참조: An image is worth 16x16 words: Transformers for image recognition at scale

def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout = 0.):
	super().__init__()
	self.layers = nn.ModuleList([])
	for _ in range(depth):
	    self.layers.append(nn.ModuleList([
		PreNorm(dim, Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout)),
		PreNorm(dim, FeedForward(dim, mlp_dim, dropout = dropout))
	    ]))

함수에 모델을 구축했다.rearrange는 블록 수로 Attention기구와FeedForwardNetwork가 블록 수를 중복했다.

def forward(self, x):
	for attn, ff in self.layers:
	    x = attn(x) + x
	    x = ff(x) + x
return x

Transformer Encoder는Residual Learning을 사용하기 때문에__init__ 함수에서 입력과 출력의 잔차를 배운다.

Vision Transformer

class ViT(nn.Module):
    def __init__(self, *, image_size, patch_size, num_classes, dim, depth, heads, mlp_dim, pool = 'cls', channels = 3, dim_head = 64, dropout = 0., emb_dropout = 0.):
        super().__init__()
        image_height, image_width = pair(image_size)
        patch_height, patch_width = pair(patch_size)

        assert image_height % patch_height == 0 and image_width % patch_width == 0, 'Image dimensions must be divisible by the patch size.'

        num_patches = (image_height // patch_height) * (image_width // patch_width)
        patch_dim = channels * patch_height * patch_width
        assert pool in {'cls', 'mean'}, 'pool type must be either cls (cls token) or mean (mean pooling)'

        self.to_patch_embedding = nn.Sequential(
            Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = patch_height, p2 = patch_width),
            nn.Linear(patch_dim, dim),
        )

        self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))
        self.cls_token = nn.Parameter(torch.randn(1, 1, dim))
        self.dropout = nn.Dropout(emb_dropout)

        self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dropout)

        self.pool = pool
        self.to_latent = nn.Identity()

        self.mlp_head = nn.Sequential(
            nn.LayerNorm(dim),
            nn.Linear(dim, num_classes)
        )

    def forward(self, img):
        x = self.to_patch_embedding(img)
        b, n, _ = x.shape

        cls_tokens = repeat(self.cls_token, '() n d -> b n d', b = b)
        x = torch.cat((cls_tokens, x), dim=1)
        x += self.pos_embedding[:, :(n + 1)]
        x = self.dropout(x)

        x = self.transformer(x)

        x = x.mean(dim = 1) if self.pool == 'mean' else x[:, 0]

        x = self.to_latent(x)
        return self.mlp_head(x)

Vision Transformer 설치입니다.

참조: An image is worth 16x16 words: Transformers for image recognition at scale
여기서는 이해하기 쉽게 아래 설정에서 설명한다.
변량
값
이미지 크기
256x256x3
그림의 한쪽
32
class 수
1000
dim
1024
depth
6
heads
16

depth

num_patches = (image_height // patch_height) * (image_width // patch_width)

원 이미지의 종횡을 패치의 종횡으로 나누어 패치의 수량을 계산합니다.image_Height 256, patchHeight는 32이기 때문에 각각 8,num입니다.patches는 64입니다.

patch_dim = channels * patch_height * patch_width

patch_dim은 3*32*32, 3072입니다.

self.to_patch_embedding = nn.Sequential(
    Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = patch_height, p2 = patch_width),
    nn.Linear(patch_dim, dim),
)

patche 이미지를 봉인한 층입니다.그림에 부합forward.
Rearrange는 이미지 데이터를 변형하여 patch로 분할합니다.
구체적으로 입력 이미지(1, 3, 256, 256)를 (1641024)로 변환합니다.

self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))

Position Embeding을 정의합니다.
논문의 공식 Z0 = [X_{class};x_{p}^1E, ...]+ E_{pos}의 E{pos}에 해당합니다.
원본 이미지의 위치를 명확하게 제시하지 않고 매개 변수로 학습한다.또한patche수에 1을 더하는 것은 공식에도 있기 때문에 처음에 부여__init__한다.이로써 사이즈는 (1651024)입니다.

self.cls_token = nn.Parameter(torch.randn(1, 1, dim))

여기 정의Linear Projection of Flattened Patches.추가class_token를 통해 이미지 식별을 배울 수 있습니다.

self.mlp_head = nn.Sequential(
    nn.LayerNorm(dim),
    nn.Linear(dim, num_classes)
)

Transformer Encoder의 출력을 받아 클래스 분류를 하는 층.그림에 부합class_token.

class_token

x = self.to_patch_embedding(img)

이미지를 패치로 분할하여 인코딩합니다.x 사이즈는(1641024)입니다.

cls_tokens = repeat(self.cls_token, '() n d -> b n d', b = b)
x = torch.cat((cls_tokens, x), dim=1)

함수 중의cls토큰과 이미지 인코딩의 비트와 일치하는 비트를 추가했습니다.그런 다음 시작 부분에 추가합니다.cls_tokens의 사이즈는 (1,11024), x의 사이즈는 (1651024)입니다.

x += self.pos_embedding[:, :(n + 1)]

Position Embeding을 입력에 추가합니다.x의 사이즈는(1651024)입니다.

x = self.transformer(x)

Transformer Encoding을 입력하고 출력을 적용합니다.x의 사이즈는(1651024)입니다.

x = x.mean(dim = 1) if self.pool == 'mean' else x[:, 0]

Transformer Encoding 출력에서 "cls"선택Token'에 대응하는 시작 부분만 사용합니다.x의 사이즈는(1024)입니다.

self.mlp_head(x)

마지막으로 Transformer의 출력을 분류층에 입력합니다.내보낸 치수는 (1000)입니다.

끝맺다

어때, 비록 Vit의 매개 변수는 매우 많지만 설치는 매우 간단하다.앞으로 Vit를 공부하는 사람들의 이해를 도울 수 있다면 정말 좋겠다.