#%%
import torch
from torch import nn
from einops import rearrange,repeat
from einops.layers.torch import Rearrange
#%%
def pair(t):
return t if isinstance(t,tuple) else (t,t)
#%%
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)
#%%
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)
#%%
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)
#%%
class Transformer(nn.Module):
def __init__(self,dim,depth,heads,dim_head,mlp_dim,dropout=0.):
super().__init__()
self.layers = nn.ModuleList([])
# 把多个encoder堆叠在一起
for _ in range(depth):
self.layers.append(nn.ModuleList([
# encoder中的 多头注意力机制 + Add & Norm
PreNorm(dim,Attention(dim,heads=heads,dim_head=dim_head,dropout=dropout)),
# Feed Forward + Add & Norm
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
#%%
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,emb_dropout=0.,dropout=0.):
super().__init__()
image_height,image_width = pair(image_size) #224x224
patch_height,patch_width = pair(patch_size) #16x16
# 图片的高和宽必须分别整除patch_height和patch_width
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 # 把patch展平后的维度
assert pool in {'cls','mean'},'pool类型必须是cls token或者平均池化'
'''模块一:把每个patch转换为embedding'''
self.to_patch_embedding = nn.Sequential(
# (样本数,通道数,高上划分的patch数,宽上划分的patch数)-> (样本数,(patch的高,patch的宽),(patch在高上的个数,patch在宽上的个数,通道道数))
Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)',p1=patch_height,p2=patch_width),
# 1.把patch拉平,变成patch_dim。 2.把patch_dim变成dim,即embedding需要的维度。
nn.Linear(patch_dim,dim),
)
'''模块二:生成位置编码'''
# dim是encoder需要的维度,num_patches+1是把cls的位置编码也给加上去了
# nn.Parameter是pytorch的一个类,用于将一个张量标记为模型的参数,训练过程中模型将更新这些参数以最小化损失函数
# torch.randn(1,2,3) 生成一个形状为(1,2,3)的符合正态分布的随机数
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)
'''模块三:transformer的encoder'''
self.transformer = Transformer(dim,depth,heads,dim_head,mlp_dim,dropout)
# 池化操作
self.pool = pool
# 不做任何操作
self.to_latent = nn.Identity()
'''模块四:mlp_head'''
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)
#%%
v = VIT(
image_size = 224,
patch_size = 16,
num_classes = 1000,# 最后cls拿出来做linear层的时候映射到多少个维度上
dim = 1024,
depth = 6, # encoder的个数
heads = 16, # 多头注意力机制的头
mlp_dim = 2048,
dropout = 0.1,
emb_dropout = 0.1
)
#%%
img = torch.randn(1,3,224,224)
preds = v(img)
preds.shape
#%%
Vision Transformer(ViT)PyTorch代码全解析(附图解)
https://blog.csdn.net/weixin_44966641/article/details/118733341