Vision Transformer（ViT）PyTorch代碼全解析

最近CV領域的Vision Transformer將在NLP領域的Transormer結果借鑒過來，屠殺了各大CV榜單。本文將根據最原始的Vision Transformer論文，及其PyTorch實現，將整個ViT的代碼做一個全面的解析。

對原Transformer還不熟悉的讀者可以看一下Attention is All You Need原文，中文講解推薦李宏毅老師的視頻 YouTube，BiliBili 個人覺得講的很明白。

話不多說，直接開始。

下圖是ViT的整體框架圖，我們在解析代碼時會參照此圖：

以下是文中給出的符號公式，也是我們解析的重要參照：
$$\mathbf{z}=[{\mathbf{x}}_{class};{\mathbf{x}}_p^1\mathbf{E},{\mathbf{x}}_p^2\mathbf{E},\dots ;{\mathbf{x}}_p^N\mathbf{E}]+{\mathbf{E}}_{pos},\mathbf{E}\in {\mathbb{R}}^{(P^2?C)\times D},{\mathbf{E}}_{pos}\in {\mathbb{R}}^{(N+1)\times D}(1)$$

$${\mathbf{z}}_{?}^{\prime }=MSA(LN({\mathbf{z}}_{??1}))+{\mathbf{z}}_{??1}(2)$$

$${\mathbf{z}}_{?}=MLP(LN({{\mathbf{z}}^{\prime }}_{?}))+{{\mathbf{z}}^{\prime }}_{?}(3)$$

$$\mathbf{y}=LN({\mathbf{z}}_L^0)(4)$$

導入需要的包

import torch from torch import nn, einsum import torch.nn.functional as Ffrom einops import rearrange, repeat from einops.layers.torch import Rearrange

都是搭建網絡時常用的PyTorch包，其中在卷積神經網絡的搭建中并不常用的einops和einsum，還不熟悉的讀者可以參考博客：einops和einsum：直接操作張量的利器。

pair函數

def pair(t):return t if isinstance(t, tuple) else (t, t)

作用是：判斷t是否是元組，如果是，直接返回t；如果不是，則將t復制為元組(t, t)再返回。
用來處理當給出的圖像尺寸或塊尺寸是int類型（如224）時，直接返回為同值元組（如(224, 224)）。

PreNorm

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

PreNorn對應框圖中最下面的黃色的Norm層。其參數dim是維度，而fn則是預先要進行的處理函數，是以下的Attention、FeedForward之一，分別對應公式（2）（3）。
$${\mathbf{z}}_{?}^{\prime }=MSA(LN({\mathbf{z}}_{??1}))+{\mathbf{z}}_{??1}(2)$$

$${\mathbf{z}}_{?}=MLP(LN({{\mathbf{z}}^{\prime }}_{?}))+{{\mathbf{z}}^{\prime }}_{?}(3)$$

FeedForward

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)

FeedForward層由全連接層，配合激活函數GELU和Dropout實現，對應框圖中藍色的MLP。參數dim和hidden_dim分別是輸入輸出的維度和中間層的維度，dropour則是dropout操作的概率參數p。

Attention

class Attention(nn.Module): def __init__(self, dim, heads=8, dim_head=64, dropout=0.):super().__init__()inner_dim = dim_head * headsproject_out = not (heads == 1 and dim_head == dim)self.heads = headsself.scale = dim_head ** -0.5self.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):b, n, _, h = *x.shape, self.headsqkv = self.to_qkv(x).chunk(3, dim=-1) # (b, n(65), dim*3) ---> 3 * (b, n, dim)q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h=h), qkv) # q, k, v (b, h, n, dim_head(64))dots = einsum('b h i d, b h j d -> b h i j', q, k) * self.scaleattn = self.attend(dots)out = einsum('b h i j, b h j d -> b h i d', attn, v)out = rearrange(out, 'b h n d -> b n (h d)')return self.to_out(out)

Attention，Transformer中的核心部件，對應框圖中的綠色的Multi-Head Attention。參數heads是多頭自注意力的頭的數目，dim_head是每個頭的維度。

本層的對應公式就是經典的Tansformer的計算公式：
$$Attention(Q,K,V)=softmax(\frac{Q{K}^T}{\sqrt{d_k}})V$$

Transformer

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) + xx = ff(x) + xreturn x

定義好幾個層之后，我們就可以構建整個Transformer Block了，即對應框圖中的整個右半部分Transformer Encoder。有了前面的鋪墊，整個Block的實現看起來非常簡潔。

參數depth是每個Transformer Block重復的次數，其他參數與上面各個層的介紹相同。

筆者也在圖中也做了標注與代碼的各部分對應。

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__()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 == 0num_patches = (image_height // patch_height) * (image_width // patch_width)patch_dim = channels * patch_height * patch_widthassert pool in {'cls', 'mean'}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)) # nn.Parameter()定義可學習參數self.dropout = nn.Dropout(emb_dropout)self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dropout)self.pool = poolself.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 c (h p1) (w p2) -> b (h w) (p1 p2 c) -> b (h w) dimb, n, _ = x.shape # b表示batchSize, n表示每個塊的空間分辨率, _表示一個塊內有多少個值cls_tokens = repeat(self.cls_token, '() n d -> b n d', b=b) # self.cls_token: (1, 1, dim) -> cls_tokens: (batchSize, 1, dim) x = torch.cat((cls_tokens, x), dim=1) # 將cls_token拼接到patch token中去 (b, 65, dim)x += self.pos_embedding[:, :(n+1)] # 加位置嵌入（直接加） (b, 65, dim)x = self.dropout(x)x = self.transformer(x) # (b, 65, dim)x = x.mean(dim=1) if self.pool == 'mean' else x[:, 0] # (b, dim)x = self.to_latent(x) # Identity (b, dim)print(x.shape)return self.mlp_head(x) # (b, num_classes)

筆者在forward()函數代碼中的注釋說明了各中間state的尺寸形狀，可供參考比對。

在 x 送入transformer之前，都是對應公式（1）的預處理操作：
$$\mathbf{z}=[{\mathbf{x}}_{class};{\mathbf{x}}_p^1\mathbf{E},{\mathbf{x}}_p^2\mathbf{E},\dots ;{\mathbf{x}}_p^N\mathbf{E}]+{\mathbf{E}}_{pos},\mathbf{E}\in {\mathbb{R}}^{(P^2?C)\times D},{\mathbf{E}}_{pos}\in {\mathbb{R}}^{(N+1)\times D}(1)$$
positional embedding和class token由nn.Parameter()定義，該函數會將送到其中的Tensor注冊到Parameters列表，隨模型一起訓練更新，對nn.Parameter()不熟悉的同學可參考博客：PyTorch中的torch.nn.Parameter() 詳解。

我們知道，transformer模型最后送到mlp中做預測的只有cls_token的輸出結果（如上圖紅框所示），而其他的圖像塊的輸出全都不要了，是由這一步實現：

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

可以看到，如果指定池化方式為'mean'的話，則會對全部token做平均池化，然后全部進行送到mlp中，但是我們可以看到，默認的self.pool='cls'，也就是說默認不會進行平均池化，而是按照ViT的設計只使用cls_token，即x[:, 0]只取第一個token（cls_token）。

最后經過mlp_head，得到各類的預測值。

筆者也簡單做了一張圖展示整個過程中的信號流，可以結合代碼中注釋的維度的變化來看：

圖中各符號含義：$H,W,C$ 分別是某一張輸入圖像的長、寬、通道數，$h,w$ 是圖塊的長、寬，如此這張圖中塊的個數就是 $\frac{H}{h}\times \frac{W}{w}$ ，用 $N_p$ 表示，$D$ 是維度數dim，$N_c$ 是類的個數。

至此，ViT模型的定義就全部完成了，在訓練腳本中實例化一個ViT模型來進行訓練即可，以下腳本可驗證ViT模型正常運作。

model_vit = ViT(image_size = 256,patch_size = 32,num_classes = 1000,dim = 1024,depth = 6,heads = 16,mlp_dim = 2048,dropout = 0.1,emb_dropout = 0.1)img = torch.randn(16, 3, 256, 256)preds = model_vit(img) print(preds.shape) # (16, 1000)

有疑惑或異議歡迎留言討論。

創作挑戰賽新人創作獎勵來咯，堅持創作打卡瓜分現金大獎

總結

以上是生活随笔為你收集整理的Vision Transformer（ViT）PyTorch代码全解析（附图解）的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。