VIT Vision Transformer | 先從PyTorch程式碼瞭解

忽逢桃林發表於2021-02-03
ORMPyTorch
  • 文章原創自：微信公眾號「機器學習煉丹術」
  • 作者：煉丹兄
  • 聯絡方式：微信cyx645016617
  • 程式碼來自github

【前言】：看程式碼的時候，也許會不理解VIT中各種元件的含義，但是這個文章的目的是瞭解其實現。在之後看論文的時候，可以做到心中有數，而不是一片茫然。

VIT類

初始化

和之前的學習一樣，從大模型類開始看起，然後一點一點看小模型類：

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__()
        assert image_size % patch_size == 0, 'Image dimensions must be divisible by the patch size.'
        num_patches = (image_size // patch_size) ** 2
        patch_dim = channels * patch_size ** 2
        assert num_patches > MIN_NUM_PATCHES, f'your number of patches ({num_patches}) is way too small for attention to be effective (at least 16). Try decreasing your patch size'
        assert pool in {'cls', 'mean'}, 'pool type must be either cls (cls token) or mean (mean pooling)'

        self.patch_size = patch_size

        self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))
        self.patch_to_embedding = nn.Linear(patch_dim, 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)
        )

在實際的呼叫中，是如下呼叫的：

model = ViT(
    dim=128,
    image_size=224,
    patch_size=32,
    num_classes=2,
    channels=3,
).to(device)

輸入引數講解：

  • image_size：圖片的大小；
  • patch_size:把圖片劃分成小的patch，小的patch的尺寸；
  • num_classes:這次分類任務的類別總數；
  • channels:輸入圖片的通道數。

VIT類中初始化的元件：

  • num_patches:一個圖片劃分成多少個patch，因為圖片224，patch32，所以劃分成7x7=49個patches；
  • patch_dim:3x32x32，理解為一個patch中的元素個數；

......這樣展示是不是非常的麻煩，還要上下來回翻看程式碼，所以我寫成註釋的形式

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.):
    # image_size=224,patch_size=32,num_classes=2,channels=3，dim=128
        super().__init__()
        assert image_size % patch_size == 0, 'Image dimensions must be divisible by the patch size.'
        # num_pathes = （224//32）**2 = 7*7=49
        num_patches = (image_size // patch_size) ** 2
        # patch_dim = 3*32*32
        patch_dim = channels * patch_size ** 2
        assert num_patches > MIN_NUM_PATCHES, f'your number of patches ({num_patches}) is way too small for attention to be effective (at least 16). Try decreasing your patch size'
        assert pool in {'cls', 'mean'}, 'pool type must be either cls (cls token) or mean (mean pooling)'
		# self.patch_size = 32
        self.patch_size = patch_size
        # self.pos_embedding是一個形狀為（1，50，128）
        self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))
        # self.patch_to_embedding是一個從3*32*32到128對映的線性層
        self.patch_to_embedding = nn.Linear(patch_dim, dim)
        # self.cls_token是一個隨機初始化的形狀為（1，1，128）這樣的變數
        self.cls_token = nn.Parameter(torch.randn(1, 1, dim))
        self.dropout = nn.Dropout(emb_dropout)
        
        # Transformer後面會講解
        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)
        )

forward

現在看VIT的推理過程：

    def forward(self, img, mask = None):
 		# p=32
        p = self.patch_size
        x = rearrange(img, 'b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = p, p2 = p)
        x = self.patch_to_embedding(x) # x.shape=[b,49,128]
        b, n, _ = x.shape # n = 49

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

        x = self.transformer(x, mask) # x.shape=[b,50,128],mask=None

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

        x = self.to_latent(x)
        return self.mlp_head(x)
  • 這裡的程式碼用到了from einops import rearrange, repeat，這個庫函式，einops是一個庫函式，是對張量進行操作的庫函式，支援pytorch，TF等。
  • einops.rearrange是把輸入的img,從[b,3,224,224]的形狀改成[b,3,7,32,7,32]的形狀，通過矩陣的轉置換成[b,7,7,32,32,3]的樣子，最後合併成[b,49,32x32x3]
  • self.patch_to_embedding,輸出的x的形狀為[b,49,128];
  • einops.repeat是把self.cls_token從[1,1,128]複製成[b,1,128]

現在，我們知道從patch到embedding是用線性層實現的。

transformer

class Transformer(nn.Module):
    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout):
        # dim=128,depth=12，heads=8，dim_head=64,mlp_dim=128
        super().__init__()
        self.layers = nn.ModuleList([])
        for _ in range(depth):
            self.layers.append(nn.ModuleList([
                Residual(PreNorm(dim, Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout))),
                Residual(PreNorm(dim, FeedForward(dim, mlp_dim, dropout = dropout)))
            ]))
    def forward(self, x, mask = None):
        for attn, ff in self.layers:
            x = attn(x, mask = mask)
            x = ff(x)
        return x
  • self.layers中包含depth組的Attention+FeedForward模組。
  • 這裡需要記得，輸入的x的尺寸為[b,50,128]

Attention

class Attention(nn.Module):
    def __init__(self, dim, heads = 8, dim_head = 64, dropout = 0.):
        super().__init__()
        inner_dim = dim_head *  heads # 64 x 8
        self.heads = heads # 8
        self.scale = dim_head ** -0.5

        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)
        self.to_out = nn.Sequential(
            nn.Linear(inner_dim, dim),
            nn.Dropout(dropout)
        )

    def forward(self, x, mask = None):
        b, n, _, h = *x.shape, self.heads # n=50,h=8
        # self.to_qkv(x)得到的尺寸為[b,50,64x8x3],然後chunk成3份
        # 也就是說，qkv是一個三元tuple,每一份都是[b,50,64x8]的大小
        qkv = self.to_qkv(x).chunk(3, dim = -1)
        # 把每一份從[b,50,64x8]變成[b,8,50,64]的形式
        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = h), qkv)
		# 這一步不太好理解，q和k都是[b,8,50,64]的形式，50理解為特徵數量，64為特徵變數
        # dots.shape=[b,8,50,50]
        dots = torch.einsum('bhid,bhjd->bhij', q, k) * self.scale
        # 不考慮mask這一塊的內容
        mask_value = -torch.finfo(dots.dtype).max

        if mask is not None:
            mask = F.pad(mask.flatten(1), (1, 0), value = True)
            assert mask.shape[-1] == dots.shape[-1], 'mask has incorrect dimensions'
            mask = mask[:, None, :] * mask[:, :, None]
            dots.masked_fill_(~mask, mask_value)
            del mask
		# 對[b,8,50,50]的最後一個維度做softmax
        attn = dots.softmax(dim=-1)

		# 這個attn就是計算出來的自注意力值，和v做點乘，out.shape=[b,8,50,64]
        out = torch.einsum('bhij,bhjd->bhid', attn, v)
        # out.shape變成[b,50,8x64]
        out = rearrange(out, 'b h n d -> b n (h d)')
        # out.shape重新變成[b,60,128]
        out =  self.to_out(out)
        return out

綜上所屬，這個attention其實就是一個自注意力模組，輸入的是[b,50,128]，返回的也是[b,50,128]。實現的過程因為使用了torch.einsum所以有些複雜，但是總的來說，和我之前講過的一篇論文"non-local"模組，是完全一樣的。torch.einsum和torch.mm原理相同，只是因為torch.mm不支援高緯度的張量做矩陣乘法。

PreNorm

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

先對輸入的x(x.shape=[b,50,128])做一個layerNormalization層歸一化，然後再放到上面的Attention模組中做自注意力。

Residual

class Residual(nn.Module):
    def __init__(self, fn):
        super().__init__()
        self.fn = fn
    def forward(self, x, **kwargs):
        return self.fn(x, **kwargs) + x

一個殘差模組罷了。

FeedForward

class FeedForward(nn.Module):
    def __init__(self, dim, hidden_dim, dropout = 0.):
    # dim=128,hidden_dim=128
        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)

就是兩個線性層，這裡有意思的是GELU()啟用函式，這個啟用函式可以直接使用torch.nn.GELU()呼叫，回頭有機會再好好講一下GELU()的原理。

transformer總結

Residual(PreNorm(dim, Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout))),
Residual(PreNorm(dim, FeedForward(dim, mlp_dim, dropout = dropout)))
  • 第一個就是，先對輸入做layerNormalization，然後放到attention得到attention的結果，然後結果和做layerNormalization之前的輸入相加做一個殘差連結；
  • 第二個就是，x->LayerNormalization->FeedForward線性層->y,然後這個y和輸入的x相加，做殘差連線。

VIT總結

回顧一下整個流程：

  • 一個圖片224x224，分成了49個32x32的patch；
  • 對這麼多的patch做embedding，成49個128向量；
  • 再拼接一個cls_tokens，變成50個128向量；
  • 再加上pos_embedding，還是50個128向量；
  • 這些向量輸入到transformer中進行自注意力的特徵提取；
  • 輸出的是50個128向量，然後對這個50個求軍職，變成一個128向量；
  • 然後線性層把128維變成2維從而完成二分類任務的transformer模型。

問題：我對NLP瞭解不深入，有沒有人可以回答一下這個問題：cls_tokens和pos_embedding的用處是什麼？

相關文章