Vit
- Transformer的多头注意力为什么要除以sqrt(d)?
- 为了归一化 因为softmax(x)如果输入值过大,有两个问题
- 计算中含有指数 exp(x)存在数值溢出
- 输入值过大,会使得梯度过小(存在梯度消失问题)
- 为什么是sqrt(d)
- 因为估算QK^T时,Q和K的向量维度为d,当d很大的时候,内积的值可能会非常大
- 如果不进行缩放,分子可能随着d的增大而增大
- 防止softmax输出过度集中在一个值上
- 为了归一化 因为softmax(x)如果输入值过大,有两个问题
ViT采用了Transformer模型中的自注意力机制来建模图像的特征, ViT模型主体的Block结构基于Transformer的Encoder结构,包含Multi-head Attention结构。
ViT的本质:将图像视为一系列的“视觉单词”或“令牌”(tokens),而不是连续的像素数组。
- Vit先将输入图像切分成多个固定大小的图像块(patches)
(注意 每个patches由多个pixels组成)
- 这些图像块被线性嵌入到固定大小的向量中,类似于NLP中的单词嵌入。
- 每个图像块都被视为一个独立的“视觉单词”或“令牌”,并视为序列化的token输入到Transformer编码器中,从而实现了图像特征的提取和分类。
- **Patch Embadding:**首先将图片分成固定大小的图像块patches,然后将其线性展开
- **Linear Projection of Flattened Patches:**在图像分割为固定大小的图像块(patches)之后,将每个图像块展平(flatten)为一维向量,并通过一个线性变化(线性投影层)将这些一维向量转化为固定维度的嵌入向量(patch Embaddings)
- **Position Embeddings:**将Position Embeddings加入图像块中,保留位置信息
- **Classification Token:**为了完成分类任务,在序列中添加块0,叫做额外的学习分类标记
- **Transformer Encoder:**由将刚刚得到的Embadding输入transformer encoder 中,每个 tansformer encoder 则由多个堆叠的层组成,每层包括多头注意力机制(MSA)和全连接的前馈神经网络组成(MLP block)
ViT的工作流程:将图像分割为固定大小的图像块(patches),将其转换为Patch Embeddings,添加位置编码信息,通过包含多头自注意力和前馈神经网络的Transformer编码器处理这些嵌入,最后利用分类标记进行图像分类等任务。
- 集合了 类别向量、图像块嵌入 和 位置编码 三者到一体的 输入嵌入向量 后,即可馈入Transformer Encoder。
- 通过不断前向通过 由 Transformer Encoder Blocks 串行堆叠构成的 Transformer Encoder,最后 提取可学习的类别嵌入向量 —— class token 对应的特征用于 图像分类。
维度变化
- Patch Embeddings
-
依次通过 Conv + Flatten + Transpose
-
输入图像 X 为 [H,W,C],通过flatten将其转化为2D的patches序列 X_p 为 [ N, P*P, C]
N = ( H ∗ W ) / P ∗ P N = (H * W) / P*P N=(H∗W)/P∗P
N即为有效的序列长度
由于Transformer在其所有层中均使用恒定的隐向量(latent vector)大小为D,将图像patches展平,并且通过可学习的线性投影层映射为D维,同时保持序列的长度不变
此时获得Patch Embaddings 为 [N, D]
总结:Patch Embaddings 本质是对每个flatten后的 patch vetor 做一个线性变换(E为( P^2 *C)✖️ D),从 P^2 *C维度 降维至D维
import torch
import torch.nn as nn
class PatchEmbed(nn.Module):
def __init__(self, image_size = 224, patch_size = 16, in_channels = 3, embed_dim = 768):
super().__init__()
# 初始化image的大小 H,W
image_size = to_2tuple(image_size)
patch_size = to_2tuple(patch_size)
num_patches = (image_size[0] // patch_size[0]) * (image_size[1] // patch_size[1])
self.image_size = image_size
self.patch_size = patch_size
self.num_patches = num_patches
# 注意重点是初始化一个可训练的线性投影层
# 获取input的embedding
# 这个线性投影层是一个二维卷积,且初始化参数中为图片的通道数,输出embedding的维度,其中卷积核大小为patch大小,步长等于卷积核,相当于卷积核将在整个输入上滑动,覆盖所有的patches,没有重叠。
self.proj = nn.Conv2d(in_channels, embed_dim, kernel_size = patch_size, stride = patch_size)
def forward(self,x):
B, C, H, W = x.shape
# 检查维度是否匹配
assert H == self.image_size[0] and W == self.image_size[1], "Input image size doesn't match model"
# D 为 embed_dim
# (B,C,H,W)->(B,D,(H // P),(W // P))
# N 为 (H // P)*(W // P)
# ->(B,D,N)->(B,N,D)
x = self.proj(x).flatten(2).transpose(1,2)
return x
- Learnable Embedding 可学习的嵌入
- 给长度为N的Patch Embaddings 追加分类向量,用于训练Transformer时学习类别信息
- 手动添加一个**可学习的嵌入向量作为用于分类的类别向量 X𝑐𝑙𝑎𝑠𝑠,**同时与其他图像块嵌入向量一起输入到 Transformer 编码器中,最后取追加的首个可学习的嵌入向量作为类别预测结果。
- 从而,最终输入 Transformer 的嵌入向量总长度为N+1。可学习嵌入在训练时随机初始化,然后通过训练得到。
- 无论是预训练还是微调,都有一个 分类头 (Classification Head) 附加在Vit编码器输出的特征之后,从而用于图像分类。在预训练时,分类头为 一个单层 MLP;在微调时,分类头为 单个线性层 (多层感知机与线性模型类似,区别在于 MLP 相对于 FC 层数增加且引入了非线性激活函数,例如 FC + GELU + FC 形式的 MLP)。
- Position Embeddings 位置嵌入
- position Embeddings也要加入patches Embeddings,主要是为了保留输入图像块之间的空间位置信息。
- 主要是由于自注意力的扰动不变性(Permutation-invariant),即打乱sequence中的tokens的顺序并不会改变其结果。
# 随机初始化分类的标记
# 初始化形状为 (1,1,D)
self.cls_token = nn.Parameter(torch.zero(1,1,embed_dim))
# 分类头
# 全连接层 其功能是将输入数据从 self.num_features 维度映射到 num_classes 维度
# 如果num_classes = 0 则不做任何变换
self.head = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity()
# forward的过程
B = x.shape[0] # x (B,C,H,W)
x = self.PatchEmbed(x)
# cls_token 为 (B,1,D) -1 表示该维度上的大小与其原本向量一致
cls_tokens = self.cls_token.expand(B,-1,-1)
# 分类编码器是按通道拼接
# x (B,N,D)
# cls_tokens (B,1,D)
# 指代在第1维度上进行拼接
# x变成 (B,N+1,D)
x = torch.cat((cls_tokens,x),dim =1)
# pos的embedding通过+获得
# 位置编码是按照元素相加
# self.pos_embed (1,N+1,D)
self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1,embed_dim)
x = x + self.pos_embed # shape = (B, N+1, D) - Python 广播机制
- Transformer 编码器
- 由交替的多头自注意力层(MSA)和多层感知机块(MLP)构成,在每一个块之前应用层归一化(Layer Norm),在每个块之后应用残差连接(Residual Connection)
- 想起来一个点:
- 自注意力:班级内部互评,QKV是相同的来源
- 交叉注意力:Q和V不同,多个任课老师给班级所有同学评分,任课老师是Q
- 多头自注意力:每个人从不同的标准来打分,然后把这个结果拼接起来(标准为权重矩阵)
- 我们最终需要的是什么,就对这个序列做加权 该序列为V
import torch
import torch.nn as nn
class Attention(nn.Module):
def __init__(self, dim, num_heads = 8, qkv_bias = False, qk_scale = None, attn_drop = 0., proj_drop = 0.):
super().__init__():
self.num_heads = num_heads # 多头自注意力所使用的数量
# 此时每个自注意力所需要处理的特征维度为
# 总维度 // 多头注意力数量
# 目的将输入特征的维度平均分配到每个注意力头上。这样,每个头可以独立地处理输入的一部分特征,最终的结果再通过某种方式(如拼接和再次线性变换)合并起来,以产生最终的输出。
# 例如dim = 512 那么 head_dim = 512//8 = 64
head_dim = dim // num_heads
# qk_scale 是一个定义的缩放因子 可以是自定义
# 一般我们用Q和K的点积除以每个头处理的特征维度的平方根
# 目的是控制梯度的规模并提高模型的稳定性
self.scale = qk_scale or head_dim ** -0.5
# qkv的生成往往是embedding输入线性变化得到的
# 因此我们定义一个线性层输入维度dim 输出则是dim*3
self.qkv = nn.Linear(dim, dim*3, bias = qkv_bias)
# 定义Dropout层目的是为了在自注意力机制的输出上应用 dropout 正则化
# 防止模型过拟合
self.attn_drop = nn.Dropout(attn_drop)
# 注意这样的变化是为了对自注意力的输出进一步进行非线性变化
# 用于将自注意力机制的输出映射回原始特征维度
self.proj = nn.Linear(dim,dim)
# 附带 dropout
self.proj_drop = nn.Dropout(proj_drop)
def forward(self,x):
B, N, D = x.shape # 这里的N其实相当于之前的N+1
# x (B,N,D)->self.qkv(x) (B,N,3*D)
# -> reshape后为(B, N, 3, 8, D//8)
# -> .permute(2, 0, 3, 1, 4) 交换维度后 (3, B, 8, N, D//8)
# 这里的8指代的自注意力头的数量
# (3, B, num_heads, N, head_dim)
# 这样每个自注意力头的qkv都可以独立处理
qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, D // self.num_heads).permute(2, 0, 3, 1, 4)
q, k, v = qkv[0], qkv[1], qkv[2]
# 注意力矩阵计算
# @ 表示矩阵乘法 N指的是序列长度
# q (B, num_heads, N, head_dim)
# k.transpose(-2,-1) 互换最后两个维度
# 这样计算得到attn为(B, num_heads, N, N)
# 其中每个元素 attn[i, j, k, l] 表示在位置 k 的查询相对于位置 l 的键的重要性。
attn = (q @ k.transpose(-2,-1)) * self.scale
# 归一化权重矩阵 生成最终的注意力矩阵
# 表示为在给定头和批次中,q对k在不同位置的相对重要性(以概率的形式)
attn = attn.softmax(dim = -1)
# self.attn_drop 会随机将 attn 中的一部分权重置为零
attn = self.attn_drop(attn)
# attn (B, num_heads, N, N)
# v (B, num_heads, N, head_dim)
# attn @ v 为 (B, num_heads, N, head_dim)
# (attn @ v).transpose(1,2) 为 (B, N, num_heads, head_dim)
# 接着将转置后的张量重新塑形,其中 D 是最终输出的通道数或特征维度
x = (attn @ v).transpose(1,2).reshape(B, N, D)
x = self.proj(x)
x = self.proj_drop(x)
return x
- FFN (Feed-forward network)
- 前馈神经网络 包含两个FC层
import torch
import torch.nn as nn
class Mlp(nn.module):
def __init__(self, in_features, hidden_features = None, act_layer = nn.GELU, drop = 0.):
super().__init__()
out_features = out_features or in_features
hidden_features = hidden_features or in_features
self.fc1 = nn.Linear(in_features, hidden_features)
self.act = act_layer() # 非线性激活函数
self.fc2 = nn.Linear(hidden_features, out_features)
self.drop = nn.Dropout(drop)
def forward(self, x):
x = self.fc1(x)
x = self.act(x)
x = self.drop(x)
x = self.fc2(x)
x = self.drop(x)
return x
- Transformer Encoder Block
!!! 因此 Transformer Encoder Block 包含一个MSA和FFN
注意 MSA和FFN二者均有跳跃连接和层归一化操作 构成对应的Block
- 在每个块前应用 层归一化 (Layer Norm),在每个块后应用 残差连接 (Residual Connection)
import torch
import torch.nn as nn
class Block(nn.Module):
def __init__(self, dim, num_heads, mlp_ratio = 4.,qkv_bias = False, qk_scale = None, drop = 0., attn_drop = 0., drop_path = 0., act_layer = nn.GELU, norm_layer = nn.LayerNorm):
super().__init__()
# 初始化多头自注意力层 MHA
self.attn = Attention(dim, num_heads = num_heads, qkv_bias = qkv_bias, qk_scale = qk_scale, attn_drop = attn_drop, proj_drop = drop)
# 接在MHA前的的归一化层
self.norm1 = norm_layer(dim)
# 实现 Transformer 模型中的一个变体,称为"Stochastic Depth",它是一种正则化技术,用于减少过拟合并提高模型的泛化能力
self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
# 接在MLP前的归一化层
self.norm2 = norm_layer(dim)
# 调整mlp层的隐藏层维度
# 通过调整 mlp_ratio,可以根据具体任务和数据调整模型的大小和复杂度
mlp_hidden_dim = int(dim * mlp_ratio)
self.mlp = Mlp(in_features = dim, hidden_features = mlp_hiddden_dim, act_layer = act_layer, drop = drop)
def forward(self, x):
# MHA + Add & Layer Norm
x = x + self.drop_path(self.attn(self.norm1(x)))
# MLP + Add & Layer Norm
x = x + self.drop_path(self.Mlp(self.norm2(x)))
return x
参考链接! 大佬很优秀
- 完整代码
import torch
import torch.nn as nn
class ViT(nn.Module):
def __init__(self, image_size=224, patch_size=16, in_channels=3, num_classes=1000, embed_dim=768, depth=12, num_heads=12):
super().__init__()
self.patch_embed = PatchEmbedding(image_size, patch_size, in_channels, embed_dim)
# 分类标记
self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
# 位置嵌入
num_patches = self.patch_embed.num_patches
self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
self.pos_drop = nn.Dropout(p=0.1)
# Transformer编码器
self.blocks = nn.ModuleList([
TransformerBlock(embed_dim, num_heads)
for _ in range(depth)
])
self.norm = nn.LayerNorm(embed_dim)
# 分类头
self.head = nn.Linear(embed_dim, num_classes)
def forward(self, x):
B = x.shape[0]
x = self.patch_embed(x)
cls_tokens = self.cls_token.expand(B, -1, -1)
x = torch.cat((cls_tokens, x), dim=1)
x = x + self.pos_embed
x = self.pos_drop(x)
for blk in self.blocks:
x = blk(x)
x = self.norm(x)
return self.head(x[:, 0])
class TransformerBlock(nn.Module):
def __init__(self, embed_dim, num_heads):
super().__init__()
self.attn = nn.MultiheadAttention(embed_dim, num_heads)
self.norm1 = nn.LayerNorm(embed_dim)
self.norm2 = nn.LayerNorm(embed_dim)
self.mlp = nn.Sequential(
nn.Linear(embed_dim, embed_dim*4),
nn.GELU(),
nn.Linear(embed_dim*4, embed_dim)
)
self.dropout = nn.Dropout(p=0.1)
def forward(self, x):
x = x + self.dropout(self.attn(x, x, x)[0])
x = self.norm1(x)
x = x + self.dropout(self.mlp(x))
x = self.norm2(x)
return x
class PatchEmbedding(nn.Module):
def __init__(self, image_size=224, patch_size=16, in_channels=3, embed_dim=768):
super().__init__()
image_size = to_2tuple(image_size)
patch_size = to_2tuple(patch_size)
num_patches = (image_size[0] // patch_size[0]) * (image_size[1] // patch_size[1])
self.image_size = image_size
self.patch_size = patch_size
self.num_patches = num_patches
self.proj = nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=patch_size)
def forward(self, x):
B, C, H, W = x.shape
assert H == self.image_size[0] and W == self.image_size[1], "Input image size doesn't match model"
x = self.proj(x).flatten(2).transpose(1, 2)
return x
def to_2tuple(x):
return (x, x) if isinstance(x, int) else x