Vision Transformer（1）：ViT源码逐行阅读解析
上图是Vision Transformer原文的模型结构展示，可以看到模型包含了几个核心模块：
Vision Transformer:
1. Embedding模块
2.Transformer Encoder模块
2.1 NormLayer ( × depth )
2.1.1 Multi-Head Attention层
关于Attention机制的详细解析
2.1.2 MLP多层感知器
3.MLP-Head 模块映射为类别
自底向上摸索是在未知中探索的不可缺少的方式，但通过摸索后，发现自顶向下能更好的阐述清楚整个逻辑。
一、ViT & Embedding
假设训练数据维度为(64, 3, 224, 224),意味着有64张三通道的224*224的图像。
设定参数dim=128意味着编码向量长度为128。
ViT中出现的PreNorm、Attention、FeedForward、Transformer后续解释
class ViT(nn.Module):
'''
:param
*: input data
image_size: 等边图像尺寸
patch_size: patch的尺寸
num_classes: 分类类别
dim: 为每一个patch编码的长度
depth: Encoder的深度，也就是连接encoder的数目
heads: 多头注意力中头的数目
mlp_dim: 多层感知器中隐含层的维度
pool: 使用cls token还是使用均值池化
channel: 图像的通道数
dim_head: 注意力机制中一个头的输入维度
dropout: NormLayer中dropout的参数比例
emb_dropout: Embedding中的dropout比例
:return 分类结果(64, 2)
'''
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_size就是每一张图像的长和宽，通过pair函数便捷明了的表现
# patch_size就是图像的每一个patch的长和宽
image_height, image_width = pair(image_size)
patch_height, patch_width = pair(patch_size)
# 保证图像可以整除为若干个patch
assert image_height % patch_height == 0 and image_width % patch_width == 0, 'Image dimensions must be divisible by the patch size.'
# 计算出每一张图片会被切割为多少个patch
# 假设输入维度(64, 3, 224, 224), num_patches = 49
num_patches = (image_height // patch_height) * (image_width // patch_width)
# 每一个patch数组大小, patch_dim = 3*32*32=3072
patch_dim = channels * patch_height * patch_width
# cls就是分类的Token， mean就是均值池化
assert pool in {'cls', 'mean'}, 'pool type must be either cls (cls token) or mean (mean pooling)'
# embeding操作：假设输入维度(64, 3, 224, 224)，那么经过Rearange层后变成了(64, 7*7=49, 32*32*3=3072)
self.to_patch_embedding = nn.Sequential(
# 将图片分割为b*h*w个三通道patch，b表示输入图像数量
Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = patch_height, p2 = patch_width),
# 经过线性全连接后，维度变成(64, 49, 128)
nn.Linear(patch_dim, dim),
)
# dim张图像，每张图像需要num_patches个向量进行编码
# 位置编码(1, 50, 128) 本应该为49，但因为cls表示类别需要增加一个
self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))
# CLS类别token,(1, 1, 128)
self.cls_token = nn.Parameter(torch.randn(1, 1, dim))
# 设置dropout
self.dropout = nn.Dropout(emb_dropout)
# 初始化Transformer
self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dropout)
# pool默认是cls进行分类
self.pool = pool
self.to_latent = nn.Identity()
# 多层感知用于将最终特征映射为2个类别
self.mlp_head = nn.Sequential(
nn.LayerNorm(dim),
nn.Linear(dim, num_classes)
)
def forward(self, img):
# 第一步，原始图像ebedding，进行了图像切割以及线性变换，变成x->(64, 49, 128)
x = self.to_patch_embedding(img)
# 得到原始图像数目和单图像的patches数量, b=64, n=49
b, n, _ = x.shape
# (1, 1, 128) -> (64, 1, 128) 为每一张图像设置一个cls的token
cls_tokens = repeat(self.cls_token, '() n d -> b n d', b = b)
# 将cls token加入到数据中 -> (64, 50, 128)
x = torch.cat((cls_tokens, x), dim=1)
# x(64, 50, 128)添加位置编码(1, 50, 128)
x += self.pos_embedding[:, :(n + 1)]
# 经过dropout层防止过拟合
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)二、Transformer
class Transformer(nn.Module):
def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout = 0.):
super().__init__()
# 设定depth个encoder相连，并添加残差结构
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):
# 每次取出包含Norm-attention和Norm-mlp这两个的ModuleList，实现残差结构
for attn, ff in self.layers:
x = attn(x) + x
x = ff(x) + x
return x1.Norm层
class PreNorm(nn.Module):
'''
:param dim 输入维度
fn 前馈网络层，选择Multi-Head Attn和MLP二者之一
'''
def __init__(self, dim, fn):
super().__init__()
# LayerNorm: ( a - mean(last 2 dim) ) / sqrt( var(last 2 dim) )
# 数据归一化的输入维度设定，以及保存前馈层
self.norm = nn.LayerNorm(dim)
self.fn = fn
# 前向传播就是将数据归一化后传递给前馈层
def forward(self, x, **kwargs):
return self.fn(self.norm(x), **kwargs)2.MLP层
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)3.Attention层
class Attention(nn.Module):
def __init__(self, dim, heads = 8, dim_head = 64, dropout = 0.):
super().__init__()
inner_dim = heads * dim_head
project_out = not (heads == 1 and dim_head == dim)
self.heads = heads
# 表示1/(sqrt(dim_head))用于消除误差，保证方差为1，避免向量内积过大导致的softmax将许多输出置0的情况
# 可以看原文《attention is all you need》中关于Scale Dot-Product Attention如何抑制内积过大
self.scale = dim_head ** -0.5
# dim = > 0 时，表示mask第d维度，对相同的第d维度，进行softmax
# dim = < 0 时，表示mask倒数第d维度，对相同的倒数第d维度，进行softmax
self.attend = nn.Softmax(dim = -1)
# 生成qkv矩阵，三个矩阵被放在一起，后续会被分开
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):
# 分割成q、k、v三个矩阵
# qkv为 inner_dim * 3，其中inner_dim = heads * dim_head
qkv = self.to_qkv(x).chunk(3, dim = -1)
# qkv的维度是(3, inner_dim = heads * dim_head)
# 'b n (h d) -> b h n d' 重新按思路分离出8个头，一共8组q,k,v矩阵
# rearrange后维度变成 (3, heads, dim, dim_head)
# 经过map后，q、k、v维度变成(1, heads, dim, dim_head)
q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), qkv)
# query * key 得到对value的注意力预测，并通过向量内积缩放防止softmax无效化部分参数
# heads * dim * dim
dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale
# 对最后一个维度进行softmax后得到预测的概率值
attn = self.attend(dots)
# 乘积得到预测结果
# out -> heads * dim * dim_head
out = torch.matmul(attn, v)
# 重组张量，将heads维度重新还原
out = rearrange(out, 'b h n d -> b n (h d)')
return self.to_out(out)4.其他部分
def pair(t):
return t if isinstance(t, tuple) else (t, t)三、MLP-Head模块
self.mlp_head = nn.Sequential(
nn.LayerNorm(dim),
nn.Linear(dim, num_classes)
)