【第三十一周】ViT 论文阅读笔记 - IT评测·应用市场-qidao123.com技术社区

class PatchEmbed(nn.Module):
"""
2D Image to Patch Embedding
"""
def __init__(self, img_size=224, patch_size=16, in_c=3, embed_dim=768, norm_layer=None):
"""
初始化 PatchEmbed 模块
Args:
img_size (int or tuple): 输入图像的尺寸，默认为 224
patch_size (int or tuple): 图像块的尺寸，默认为 16
in_c (int): 输入图像的通道数，默认为 3
embed_dim (int): 嵌入维度，默认为 768
norm_layer (nn.Module): 归一化层，默认为 None
"""
super().__init__()
# 将图像尺寸转换为元组形式
img_size = (img_size, img_size)
# 将图像块尺寸转换为元组形式
patch_size = (patch_size, patch_size)
self.img_size = img_size
self.patch_size = patch_size
# 计算网格尺寸
self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])
# 计算图像块的数量
self.num_patches = self.grid_size[0] * self.grid_size[1]
# 定义卷积层，用于将图像分割成图像块并进行嵌入
self.proj = nn.Conv2d(in_c, embed_dim, kernel_size=patch_size, stride=patch_size)
# 定义归一化层，如果提供了则使用，否则使用恒等映射
self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
def forward(self, x):
"""
前向传播
Args:
x (torch.Tensor): 输入图像张量，形状为 [B, C, H, W]
Returns:
torch.Tensor: 处理后的张量，形状为 [B, num_patches, embed_dim]
"""
# 获取输入图像的形状
B, C, H, W = x.shape
# 检查输入图像的尺寸是否与模型设置的尺寸一致
assert H == self.img_size[0] and W == self.img_size[1], \
f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
# 通过卷积层将图像分割成图像块并进行嵌入
# flatten: [B, C, H, W] -> [B, C, HW]
# transpose: [B, C, HW] -> [B, HW, C]
x = self.proj(x).flatten(2).transpose(1, 2)
# 进行归一化处理
x = self.norm(x)
return x

复制代码

class Attention(nn.Module):
def __init__(self,
dim, # 输入token的dim
num_heads=8,
qkv_bias=False,
qk_scale=None,
attn_drop_ratio=0.,
proj_drop_ratio=0.):
"""
初始化 Attention 模块
Args:
dim (int): 输入 token 的维度
num_heads (int): 注意力头的数量，默认为 8
qkv_bias (bool): 是否使用偏置项，默认为 False
qk_scale (float): 缩放因子，默认为 None
attn_drop_ratio (float): 注意力矩阵的丢弃概率，默认为 0.
proj_drop_ratio (float): 投影层的丢弃概率，默认为 0.
"""
super(Attention, self).__init__()
self.num_heads = num_heads
# 计算每个注意力头的维度
head_dim = dim // num_heads
# 计算缩放因子，如果未提供则使用默认值
self.scale = qk_scale or head_dim ** -0.5
# 定义线性层，用于生成查询（Q）、键（K）和值（V）
self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
# 定义注意力矩阵的丢弃层
self.attn_drop = nn.Dropout(attn_drop_ratio)
# 定义投影层
self.proj = nn.Linear(dim, dim)
# 定义投影层的丢弃层
self.proj_drop = nn.Dropout(proj_drop_ratio)
def forward(self, x):
"""
前向传播
Args:
x (torch.Tensor): 输入张量，形状为 [batch_size, num_patches + 1, total_embed_dim]
Returns:
torch.Tensor: 处理后的张量，形状为 [batch_size, num_patches + 1, total_embed_dim]
"""
# 获取输入张量的形状
B, N, C = x.shape
# 通过线性层生成查询（Q）、键（K）和值（V）
# qkv(): -> [batch_size, num_patches + 1, 3 * total_embed_dim]
# reshape: -> [batch_size, num_patches + 1, 3, num_heads, embed_dim_per_head]
# permute: -> [3, batch_size, num_heads, num_patches + 1, embed_dim_per_head]
qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
# 分离查询（Q）、键（K）和值（V）
# [batch_size, num_heads, num_patches + 1, embed_dim_per_head]
q, k, v = qkv[0], qkv[1], qkv[2] # make torchscript happy (cannot use tensor as tuple)
# 计算注意力分数
# transpose: -> [batch_size, num_heads, embed_dim_per_head, num_patches + 1]
# @: multiply -> [batch_size, num_heads, num_patches + 1, num_patches + 1]
attn = (q @ k.transpose(-2, -1)) * self.scale
# 对注意力分数进行 softmax 操作，得到注意力矩阵
attn = attn.softmax(dim=-1)
# 对注意力矩阵进行丢弃操作
attn = self.attn_drop(attn)
# 根据注意力矩阵对值（V）进行加权求和
# @: multiply -> [batch_size, num_heads, num_patches + 1, embed_dim_per_head]
# transpose: -> [batch_size, num_patches + 1, num_heads, embed_dim_per_head]
# reshape: -> [batch_size, num_patches + 1, total_embed_dim]
x = (attn @ v).transpose(1, 2).reshape(B, N, C)
# 通过投影层进行线性变换
x = self.proj(x)
# 对投影层的输出进行丢弃操作
x = self.proj_drop(x)
return x

复制代码

class Mlp(nn.Module):
"""
MLP as used in Vision Transformer, MLP-Mixer and related networks
"""
def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
"""
初始化 MLP 模块
Args:
in_features (int): 输入特征的维度
hidden_features (int): 隐藏层特征的维度，默认为 None
out_features (int): 输出特征的维度，默认为 None
act_layer (nn.Module): 激活函数层，默认为 nn.GELU
drop (float): 丢弃概率，默认为 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):
"""
前向传播
Args:
x (torch.Tensor): 输入张量
Returns:
torch.Tensor: 处理后的张量
"""
# 通过第一个全连接层
x = self.fc1(x)
# 通过激活函数层
x = self.act(x)
# 进行丢弃操作
x = self.drop(x)
# 通过第二个全连接层
x = self.fc2(x)
# 进行丢弃操作
x = self.drop(x)
return x

复制代码

class Block(nn.Module):
def __init__(self,
dim,
num_heads,
mlp_ratio=4.,
qkv_bias=False,
qk_scale=None,
drop_ratio=0.,
attn_drop_ratio=0.,
drop_path_ratio=0.,
act_layer=nn.GELU,
norm_layer=nn.LayerNorm):
"""
初始化 Block 模块
Args:
dim (int): 输入特征的维度
num_heads (int): 注意力头的数量
mlp_ratio (float): MLP 隐藏层维度与输入维度的比例，默认为 4.
qkv_bias (bool): 是否使用偏置项，默认为 False
qk_scale (float): 缩放因子，默认为 None
drop_ratio (float): 丢弃概率，默认为 0.
attn_drop_ratio (float): 注意力矩阵的丢弃概率，默认为 0.
drop_path_ratio (float): 随机深度丢弃概率，默认为 0.
act_layer (nn.Module): 激活函数层，默认为 nn.GELU
norm_layer (nn.Module): 归一化层，默认为 nn.LayerNorm
"""
super(Block, self).__init__()
# 定义第一个归一化层
self.norm1 = norm_layer(dim)
# 定义注意力模块
self.attn = Attention(dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale,
attn_drop_ratio=attn_drop_ratio, proj_drop_ratio=drop_ratio)
# 定义随机深度丢弃层，如果丢弃概率大于 0 则使用，否则使用恒等映射
self.drop_path = DropPath(drop_path_ratio) if drop_path_ratio > 0. else nn.Identity()
# 定义第二个归一化层
self.norm2 = norm_layer(dim)
# 计算 MLP 隐藏层的维度
mlp_hidden_dim = int(dim * mlp_ratio)
# 定义 MLP 模块
self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop_ratio)
def forward(self, x):
"""
前向传播
Args:
x (torch.Tensor): 输入张量
Returns:
torch.Tensor: 处理后的张量
"""
# 先进行归一化，再通过注意力模块，最后加上随机深度丢弃和残差连接
x = x + self.drop_path(self.attn(self.norm1(x)))
# 先进行归一化，再通过 MLP 模块，最后加上随机深度丢弃和残差连接
x = x + self.drop_path(self.mlp(self.norm2(x)))
return x

复制代码

class VisionTransformer(nn.Module):
def __init__(self, img_size=224, patch_size=16, in_c=3, num_classes=1000,
embed_dim=768, depth=12, num_heads=12, mlp_ratio=4.0, qkv_bias=True,
qk_scale=None, representation_size=None, distilled=False, drop_ratio=0.,
attn_drop_ratio=0., drop_path_ratio=0., embed_layer=PatchEmbed, norm_layer=None,
act_layer=None):
"""
Args:
img_size (int, tuple): 输入图像的尺寸，如果是整数则表示正方形图像的边长
patch_size (int, tuple): 图像分块的尺寸，如果是整数则表示正方形分块的边长
in_c (int): 输入图像的通道数，通常彩色图像为3
num_classes (int): 分类任务的类别数
embed_dim (int): 嵌入向量的维度
depth (int): 变压器（Transformer）的层数
num_heads (int): 多头注意力机制中的头数
mlp_ratio (int): 多层感知机（MLP）隐藏层维度与嵌入维度的比例
qkv_bias (bool): 是否在查询（Q）、键（K）、值（V）的线性变换中使用偏置
qk_scale (float): 自定义的查询和键的缩放因子，如果未设置则使用默认值
representation_size (Optional[int]): 如果设置，则启用并将表示层（预对数层）的维度设置为该值
distilled (bool): 模型是否包含蒸馏令牌和头，如DeiT模型
drop_ratio (float): 随机失活（Dropout）的概率
attn_drop_ratio (float): 注意力机制中的随机失活概率
drop_path_ratio (float): 随机深度（Stochastic Depth）的概率
embed_layer (nn.Module): 用于图像分块嵌入的层
norm_layer: (nn.Module): 归一化层
"""
super(VisionTransformer, self).__init__()
# 分类任务的类别数
self.num_classes = num_classes
# 特征维度，与嵌入维度保持一致，便于与其他模型统一接口
self.num_features = self.embed_dim = embed_dim
# 令牌数量，如果使用蒸馏则为2（分类令牌和蒸馏令牌），否则为1（分类令牌）
self.num_tokens = 2 if distilled else 1
# 如果未提供归一化层，则使用默认的LayerNorm层，设置eps为1e-6
norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)
# 如果未提供激活函数层，则使用默认的GELU激活函数
act_layer = act_layer or nn.GELU
# 图像分块嵌入层，将输入图像分割成多个分块并进行嵌入
self.patch_embed = embed_layer(img_size=img_size, patch_size=patch_size, in_c=in_c, embed_dim=embed_dim)
# 分块的数量
num_patches = self.patch_embed.num_patches
# 分类令牌，可学习的参数，形状为 [1, 1, embed_dim]
self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
# 蒸馏令牌，如果使用蒸馏则为可学习的参数，形状为 [1, 1, embed_dim]，否则为None
self.dist_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) if distilled else None
# 位置嵌入，可学习的参数，形状为 [1, num_patches + num_tokens, embed_dim]
self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim))
# 位置嵌入后的随机失活层
self.pos_drop = nn.Dropout(p=drop_ratio)
# 随机深度衰减规则，从0到drop_path_ratio线性插值生成depth个值
dpr = [x.item() for x in torch.linspace(0, drop_path_ratio, depth)]
# 变压器块序列，包含多个Block层
self.blocks = nn.Sequential(*[
Block(dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
drop_ratio=drop_ratio, attn_drop_ratio=attn_drop_ratio, drop_path_ratio=dpr[i],
norm_layer=norm_layer, act_layer=act_layer)
for i in range(depth)
])
# 归一化层，用于对变压器块的输出进行归一化
self.norm = norm_layer(embed_dim)
# 表示层（预对数层）
if representation_size and not distilled:
# 如果设置了表示层维度且不使用蒸馏，则启用表示层
self.has_logits = True
# 更新特征维度为表示层维度
self.num_features = representation_size
# 表示层，包含一个线性层和一个Tanh激活函数
self.pre_logits = nn.Sequential(OrderedDict([
("fc", nn.Linear(embed_dim, representation_size)),
("act", nn.Tanh())
]))
else:
# 否则不启用表示层，使用恒等映射
self.has_logits = False
self.pre_logits = nn.Identity()
# 分类头
self.head = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity()
# 蒸馏头，如果使用蒸馏则为线性层，否则为None
self.head_dist = None
if distilled:
self.head_dist = nn.Linear(self.embed_dim, self.num_classes) if num_classes > 0 else nn.Identity()
# 权重初始化
# 位置嵌入的权重使用截断正态分布初始化，标准差为0.02
nn.init.trunc_normal_(self.pos_embed, std=0.02)
if self.dist_token is not None:
# 蒸馏令牌的权重使用截断正态分布初始化，标准差为0.02
nn.init.trunc_normal_(self.dist_token, std=0.02)
# 分类令牌的权重使用截断正态分布初始化，标准差为0.02
nn.init.trunc_normal_(self.cls_token, std=0.02)
# 应用自定义的权重初始化函数
self.apply(_init_vit_weights)
def forward_features(self, x):
"""
前向传播特征提取部分
Args:
x (torch.Tensor): 输入图像，形状为 [B, C, H, W]
Returns:
torch.Tensor: 特征向量，如果使用蒸馏则返回分类令牌和蒸馏令牌的特征向量
"""
# [B, C, H, W] -> [B, num_patches, embed_dim]
x = self.patch_embed(x) # [B, 196, 768]
# [1, 1, 768] -> [B, 1, 768]
cls_token = self.cls_token.expand(x.shape[0], -1, -1)
if self.dist_token is None:
# 如果不使用蒸馏，将分类令牌和分块嵌入拼接
x = torch.cat((cls_token, x), dim=1) # [B, 197, 768]
else:
# 如果使用蒸馏，将分类令牌、蒸馏令牌和分块嵌入拼接
x = torch.cat((cls_token, self.dist_token.expand(x.shape[0], -1, -1), x), dim=1)
# 位置嵌入并进行随机失活
x = self.pos_drop(x + self.pos_embed)
# 通过变压器块序列
x = self.blocks(x)
# 归一化
x = self.norm(x)
if self.dist_token is None:
# 如果不使用蒸馏，返回分类令牌的特征向量
return self.pre_logits(x[:, 0])
else:
# 如果使用蒸馏，返回分类令牌和蒸馏令牌的特征向量
return x[:, 0], x[:, 1]
def forward(self, x):
"""
前向传播函数
Args:
x (torch.Tensor): 输入图像，形状为 [B, C, H, W]
Returns:
torch.Tensor: 分类结果，如果使用蒸馏则返回分类结果和蒸馏结果的平均值
"""
# 提取特征
x = self.forward_features(x)
if self.head_dist is not None:
# 如果使用蒸馏，分别通过分类头和蒸馏头
x, x_dist = self.head(x[0]), self.head_dist(x[1])
if self.training and not torch.jit.is_scripting():
# 训练时返回分类结果和蒸馏结果
return x, x_dist
else:
# 推理时返回分类结果和蒸馏结果的平均值
return (x + x_dist) / 2
else:
# 如果不使用蒸馏，通过分类头得到分类结果
x = self.head(x)
return x

复制代码