萌新的卷積神經(jīng)網(wǎng)絡學習筆記--pytorch搭建Vit Transformer
一些函數(shù)的使用
tensor.chunk(劃分為多少份，dim= 按哪個維度劃分)
from einops import rearrange
from einops.layers.torch import Rearrange
類似reshape( )，view( ), 但更直觀，前者為方法，后者為類
from einops import?repeat
復制矩陣里的元素生成新矩陣，且屬于淺拷貝
einsum()
求兩個矩陣的點積
patch + position embedding
Position Embedding 中位置編碼，cls token的解讀: http://t.csdn.cn/twonr
# patch + position embedding
class Patch_embed(nn.Module):
? ?def __init__(self, img_size, patch_size,dim,P):
? ? ? ?super(Patch_embed,self).__init__()
? ? ? ?
? ? ? ?self.unfold = nn.Sequential(nn.Unfold(kernel_size = (patch_size,patch_size),
? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ?dilation=(1,1),
? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ?padding=(0,0),
? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ?stride=(patch_size,patch_size)),
? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ?Rearrange("b i j -> b j i"),
? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ?nn.Linear(patch_size**2*3,dim)
? ? ? ?)
? ? ? ?self.dropout = nn.Dropout(p=P)
? ? ? ?n = (img_size//patch_size)**2
? ? ? ?
? ? ? ?# nn.Parameter()將變量轉化為可學習的參數(shù)
? ? ? ?# 可學習的參數(shù)做位置編碼,編碼過程直接相加
? ? ? ?self.pos = nn.Parameter(torch.randn(1, n+1 ?,dim))# [1,197,dim]
? ? ? ?
? ? ? ?# 其他部分都是圖像本身各個patch間在計算，單獨在首行加一條可學習的向量（cls token）做為輸出預測的依據(jù)
? ? ? ?self.cls = nn.Parameter(torch.randn(1,1, dim))# [1,1,dim]
? ? ? ?
? ?def forward(self,x): # 輸入[1,3,224,224]
? ? ? ?
? ? ? ?# patch_size=(16,16) 224/16=14
? ? ? ?# [1,16*16*3,14*14] -> [1,14*14,16*16*3] ->[1,14*14,dim] ?
? ? ? ?# [1,768, 196] -> [1,196,768] ->[1,196,dim]
? ? ? ?x = self.unfold(x)
? ? ? ?b, _, _ = x.shape
? ? ? ?cls_token = repeat(self.cls, "() n d -> b n d",b=b) # ->[batch_size, 1, dim]
? ? ? ?# cls token加在首行
? ? ? ?x = torch.cat((cls_token,x),dim=-2) # [b,196,dim]->[b,197,dim]
? ? ? ?#position embedding
? ? ? ?x += self.pos
? ? ?
? ? ? ?x=self.dropout(x)
? ? ? ?
? ? ? ?return x主要部分：
MSA（Multi-head Attention部分）
多頭自注意力機制（大概是這樣吧）
class MSA(nn.Module):
? ?'''多頭自注意力機制，輸入輸出維度不變'''
? ?def __init__(self,dim_in, # 輸入向量（token）的維度
? ? ? ? ? ? ? ? heads, # 頭數(shù)
? ? ? ? ? ? ? ? head_dim, # 每個頭的維度
? ? ? ? ? ? ? ? P=0.
? ? ? ? ? ? ? ? ):
? ? ? ?super().__init__()
? ? ? ?self.heads = heads
? ? ? ?dim_inner = heads * head_dim
? ? ? ?self.scale = head_dim ** (-0.5) # 1/(dk**0.5)
? ? ? ?self.to_qkv = nn.Linear(dim_in, dim_inner*3) # 生成Q，K，V三個矩陣，所以乘三
? ? ? ?self.softmax = nn.Softmax(dim=-1)
? ? ?
? ? ? ?# 統(tǒng)一輸出向量的維度， 與輸入一致
? ? ? ?# 判斷要輸出的向量維度是否和輸入一致，當頭數(shù)為1并且每個頭的維度與輸入相同時 直接輸出
? ? ? ?project_out = not (heads == 1 and head_dim == dim_in)
? ? ? ?self.to_out = nn.Sequential(
? ? ? ? ? ?nn.Linear(dim_inner, dim_in),
? ? ? ? ? ?nn.Dropout(P)
? ? ? ?) if project_out else nn.Identity()
? ?
? ?def forward(self,x):
? ? ? ?# x.shape:[b, n, l],
? ? ? ?# 表示同時傳入b個數(shù)據(jù)，一個數(shù)據(jù)包含n個長度為l的向量（token）
? ? ? ?b, n, _ = *x.shape,
? ? ? ?h = self.heads
? ? ? ?
? ? ? ?# 通過一個全鏈層獲得Q,K,V矩陣,
? ? ? ?QKV = self.to_qkv(x).chunk(3,dim=-1) ?# chunk(3,dim=-1)按最后一個維度分為三份，打包成元組
? ? ? ?# map(func,iterable)
? ? ? ?Q, K, V = map(lambda t: rearrange(t,"b n (heads head_dim) -> b heads n head_dim", heads=h), QKV)
? ? ? ?
? ? ? ?# attention(q,k,v) = softmax((q*k)/dk**0.5)*v
? ? ? ?# einsum()里維度的描述只能用字母a~z, 或A~Z，不然會報錯
? ? ? ?# ? ? ? ? ? ?b heads n head_dim, b heads n head_dim -> b heads n n # -> b heads n n 的話有兩個n會報錯
? ? ? ?QK = einsum("b h i d, b h j d -> b h i j", Q, K)*self.scale # einsum('i j, k j->i k', A, B),求兩矩陣的點積
? ? ? ?
? ? ? ?A = self.softmax(QK)
? ? ? ?
? ? ? ?# ? ? ? ? ? ? b heads n n, b heads n head_dim -> b heads n head_dim
? ? ? ?att = einsum("b h i j, b h j d -> b h i d", A, V) # einsum('i j, j k->i k', A, B)
? ? ? ?
? ? ? ?# ? ? ? ? ? ? ? ?b heads n head_dim -> b n ?heads*head_dim (即dim_inner)
? ? ? ?out = rearrange(att, "b heads n head_dim -> b n (heads head_dim)")
? ? ? ?
? ? ? ?out = self.to_out(out) # 統(tǒng)一輸出的向量（token）維度與輸入維度相同
? ? ? ?
? ? ? ?return outMLP：多層感知機（即全連接層）
class MLP(nn.Module):
? ?'''通過兩層全連接層 向量維度 放大再縮小后輸出'''
? ?def __init__(self, dim_in, dim_out, P=0.3):
? ? ? ?super(MLP,self).__init__()
? ? ? ?
? ? ? ?self.mlp =nn.Sequential(
? ? ? ? ? ?nn.Linear(dim_in,dim_out),
? ? ? ? ? ?nn.GELU(),
? ? ? ? ? ?nn.Dropout(p=P),
? ? ? ? ? ?nn.Linear(dim_out,dim_in),
? ? ? ? ? ?nn.Dropout(p=P) ? ? ? ?
? ? ? ?)
? ?
? ?def forward(self,x):
? ? ? ?
? ? ? ?x = self.mlp(x)
? ? ? ?
? ? ? ?return x ?
Transformer Encoder
class Transformer_encoder(nn.Module):
? ?def __init__(self, dim_in, #輸入向量的維度
? ? ? ? ? ? ? ? heads, # 多頭自注意力機制頭數(shù)
? ? ? ? ? ? ? ? head_dim, # 每個頭的維度
? ? ? ? ? ? ? ? P, # dropout的概率
? ? ? ? ? ? ? ? mlp_dim # MLP里隱藏層的神經(jīng)元個數(shù)
? ? ? ? ? ? ? ? ) :
? ? ? ?super().__init__()
? ? ? ?
? ? ? ?self.Norm = nn.LayerNorm(dim_in) # Norm
? ? ? ?self.MSA = MSA(dim_in,heads=heads, head_dim=head_dim,P=P) # MSA
? ? ? ?self.MLP = MLP(dim_in=dim_in,mlp_dim=mlp_dim,P=P) # MLP
? ? ? ?
? ?def forward(self,x):
? ? ? ?
? ? ? ?x1 = self.Norm(x)
? ? ? ?x2 = self.MSA(x1)
? ? ? ?x3 = x + x2
? ? ? ?x4 = self.Norm(x3)
? ? ? ?x5 = self.MLP(x4)
? ? ? ?out = x3 + x5
? ? ? ?
? ? ? ?return outVIT
class vit(nn.Module):
? ?def __init__(self, img_size, # 圖片尺寸
? ? ? ? ? ? ? ? depth, # 多少層transformer encoder
? ? ? ? ? ? ? ? dim_in, # 輸入向量的維度
? ? ? ? ? ? ? ? patch_size, #分塊的大小
? ? ? ? ? ? ? ? heads, # 多頭自注意力機制頭數(shù)
? ? ? ? ? ? ? ? head_dim, # 每個頭的維度P, # dropout的概率
? ? ? ? ? ? ? ? mlp_dim, # MLP里隱藏層的神經(jīng)元個數(shù)
? ? ? ? ? ? ? ? P,
? ? ? ? ? ? ? ? class_num
? ? ? ? ? ? ? ? ):
? ? ? ?super(vit,self).__init__()
? ? ? ?self.patch_embed = Patch_embed(img_size,patch_size,dim_in,P)
? ? ? ?self.layers = nn.ModuleList()
? ? ? ?self.MLP_Head = nn.Sequential(nn.LayerNorm(dim_in),
? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ?nn.Linear(dim_in, class_num),
? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ?nn.Softmax(dim=-1)
? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ?)
? ? ? ?
? ? ? ?for i in range(depth):
? ? ? ? ? ?self.layers.append(Transformer_encoder(dim_in,heads,head_dim,mlp_dim,P))
? ? ? ?
? ?def forward(self,x): # [1, 3, 224, 224]
? ? ? ?x = self.patch_embed(x) # [[1, 197, 1000]]
? ? ? ?
? ? ? ?# 通過depth層 transformer encoder
? ? ? ?for transformer_encoder in self.layers:
? ? ? ? ? ?x = transformer_encoder(x)
? ? ? ?
? ? ? ?x = x[:,0] # 取首行的cls token
? ? ? ?
? ? ? ?x = self.MLP_Head(x) # 經(jīng)過MLP Head后輸出類別
? ? ? ?return x
test
if __name__ == "__main__":
? ?x = torch.rand(1,3,224,224)
? ?model = vit(img_size=224, # 圖片尺寸
? ? ? ? ? ? ? ? depth=3, # 多少層transformer encoder
? ? ? ? ? ? ? ? dim_in=1000, # patch_embed后的向量維度
? ? ? ? ? ? ? ? patch_size=16, #分塊的大小
? ? ? ? ? ? ? ? heads=3, # 多頭自注意力機制頭數(shù)
? ? ? ? ? ? ? ? head_dim=2000, # 每個頭的維度P
? ? ? ? ? ? ? ? mlp_dim=2000, # MLP里隱藏層的神經(jīng)元個數(shù)
? ? ? ? ? ? ? ? P=0.3,# dropout的概率即全連接層的神經(jīng)元不參與訓練的概率（防過擬合）
? ? ? ? ? ? ? ? class_num=2 # 最后輸出幾維向量（多少類）
? ? ? ? ? ? ? ? )
? ?y = model(x)
? ?
? ?print(f"x.shape: {x.shape}")
? ?print(f"y.shape: {y.shape}")
? ?print(f"y: {y}")
? ?print(f"y.sum():{y.sum()}")輸出：
不知道對不對，至少能跑通！
參考：
https://zhuanlan.zhihu.com/p/445122996
https://blog.csdn.net/weixin_44966641/article/details/118733341
