视觉常用Backbone大全：VisionTransformer（ViT）

视觉常用Backbone大全

今天介绍的主干网络模型叫VisionTransformer，是一种将 Transformer 架构应用于计算机视觉任务的模型，通过将图像进行切块，将图片转变为self-attention认识的token输入到Transformer模块中，实现了Transformer架构在视觉领域的应用；

一、模型介绍

Transformer 最初是由 Vaswani 等人在 2017 年的论文《Attention is All You Need》中提出，主要用于自然语言处理任务。2020 年，Google Research 在论文《An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale》中首次将 Transformer 应用于图像识别任务，提出了 ViT 模型。

上图左边就是ViT模型的整体架构，模型整体由三部分构成，分别是Linear Projection、Transformer Encoder、MLP Head；接下来就对这几个模块做进一步的分析；

二、模块分析

2.1 Linear Projection

在这个模块里，模型对数据进行了切分和编码操作，具体操作如下：

先说切分，如上图左下角所示，将一张图片切分成9块那样，假设原图像的尺寸为224*224*3，我们想要使得切出来的小块的尺寸为16*16，那么我们就需要将原图分成（224 / 16 )^2 = 196块，每一块的尺寸为16*16*3，这个操作叫做patch，在实际代码中，patch的裁剪是用一个patch_size大小的卷积同时以patch_size的步长进行卷积实现的；

在NLP中，一段文本在输入到transformer模块之前需要将文本送入一个可训练的encode模型进行编码，模型会将文本进行切分，再将切分后的文本转换成token张量，这样一个token张量就是一个一维的向量；

在图像中也是一样，我们将图像进行切分，但是切分后的小图像块还是三维数据，所以我们还需要将这样的每一个小图像块进行维度转换，将其变为一维的向量将其视为一个token，这样转换后的小图像块尺寸就变为了（1，768），我们一共有196个这样的图像块，所以转换维度之后的图像数据维度为（196，768）；

在转换为token张量之后，我们还需要对其添加位置信息，位置信息张量的维度与现有token维度相同，通过add方式（对应元素相加）进行融合；

最后，输入tensor还需要有一个class token（分类层），数据格式和其他token一样为（1，768），与位置编码的融合方式不一样，这里做的是Concat（维度上的拼接），这样做是因为分类信息是在后面需要取出来单独做预测的，所以不能以Add方式融合，shape也就从（196, 768）变为（197, 768）；

至此，Linear Projection部分就全部完成，token将被送入Transformer Encoder模块；

2.2 Transformer Encoder

在整个Transformer Encoder中，实际上就是如上图的几个encoder block模块的堆叠，我们可以看到在encoder block模块中主要包含Layer Normalization、Multi-Head Attention、DropOut/DropPath和MLP四部分；

2.2.1 Layer Normalization

Layer Normalization的作用和BatchNorm是同样的作用，都是将数据进行正则化处理，使得模型可以加速收敛，区别在于两者对数据处理的维度不同，BN针对的是批量数据中的某个维度进行操作，而LN则是针对某个样本的所有维度进行操作；

2.2.2 Multi-Head Attention

这里的多头自注意力机制应用的就是Transformer的Multi-Head Attention，这里就简单说一下self-attention以及Multi-Head Attention；

对于self-attention机制，它接收的是token张量，每一个token张量进入self-attention后，都会与可训练权重 $w_{q}$ 、 $w_{k}$ 、 $w_{v}$ 进行矩阵运算，生成对应的Q、K、V张量，然后不同token的Q和K进行运算（有公式）生成紧密权重 $\alpha$ 张量，最后再将 $\alpha$ 与V进行运算，生成新的token张量，这就是自注意力机制的计算流程；

而对于Multi-Head Attention，它与self-attention的区别在于同一个token可以生成多个Q、K、V张量对，中间的运算流程都是相同的，最后会输出多个新的token张量，再将其加权合并为一个token进行输出；

2.2.3 MLP

这个就是一个简单的前馈神经网络，通过全连接层、激活层、dropout层的串联对token进一步做特征提取和特征融合的操作；

2.2.4 Transformer Encoder代码

# transformer编码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):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)# NOTE: drop path for stochastic depth, we shall see if this is better than dropout hereself.drop_path = DropPath(drop_path_ratio) if drop_path_ratio > 0. else nn.Identity()self.norm2 = norm_layer(dim)mlp_hidden_dim = int(dim * mlp_ratio)self.mlp = MLP(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop_ratio)def forward(self, x):x = x + self.drop_path(self.attn(self.norm1(x)))x = x + self.drop_path(self.mlp(self.norm2(x)))return x

三、拓扑结构

这里以最常使用的ViT-B/16为例，看一下它的拓扑结构以及代码实现；

如上图所示就是ViT-B/16模型的拓扑结构，从结构上可以看出ViT-B/16模型输入尺寸为224*224，patch_siae=16*16，进过编码后输入到由12个多头自注意力机制的transformer encoder堆叠起来的网络中，最后通过一个MLP head进行分类；

四、代码实现

下面是ViT-B/16模型基于pytorch的实现：

# ViT-B/16import torch
import torch.nn as nn
import torch.nn.functional as Fclass PatchEmbedding(nn.Module):def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):super(PatchEmbedding, self).__init__()self.img_size = img_sizeself.patch_size = patch_sizeself.num_patches = (img_size // patch_size) ** 2self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)def forward(self, x):x = self.proj(x).flatten(2).transpose(1, 2)return xclass PositionalEncoding(nn.Module):def __init__(self, embed_dim, dropout=0.1, max_len=5000):super(PositionalEncoding, self).__init__()self.dropout = nn.Dropout(p=dropout)pe = torch.zeros(max_len, embed_dim)position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)div_term = torch.exp(torch.arange(0, embed_dim, 2).float() * (-math.log(10000.0) / embed_dim))pe[:, 0::2] = torch.sin(position * div_term)pe[:, 1::2] = torch.cos(position * div_term)pe = pe.unsqueeze(0).transpose(0, 1)self.register_buffer('pe', pe)def forward(self, x):x = x + self.pe[:x.size(0), :]return self.dropout(x)class TransformerEncoderLayer(nn.Module):def __init__(self, embed_dim, num_heads, mlp_dim, dropout=0.1):super(TransformerEncoderLayer, self).__init__()self.self_attn = nn.MultiheadAttention(embed_dim, num_heads, dropout=dropout)self.linear1 = nn.Linear(embed_dim, mlp_dim)self.dropout = nn.Dropout(dropout)self.linear2 = nn.Linear(mlp_dim, embed_dim)self.norm1 = nn.LayerNorm(embed_dim)self.norm2 = nn.LayerNorm(embed_dim)self.dropout1 = nn.Dropout(dropout)self.dropout2 = nn.Dropout(dropout)self.dropout3 = nn.Dropout(dropout)def forward(self, src):src2, _ = self.self_attn(src, src, src)src = src + self.dropout1(src2)src = self.norm1(src)src2 = self.linear2(self.dropout(F.relu(self.linear1(src))))src = src + self.dropout2(src2)src = self.norm2(src)return srcclass VisionTransformer(nn.Module):def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768, depth=12, num_heads=12, mlp_dim=3072, num_classes=1000):super(VisionTransformer, self).__init__()self.patch_embed = PatchEmbedding(img_size, patch_size, in_chans, embed_dim)self.pos_embed = PositionalEncoding(embed_dim, max_len=self.patch_embed.num_patches + 1)self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))self.transformer = nn.ModuleList([TransformerEncoderLayer(embed_dim, num_heads, mlp_dim) 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 = self.pos_embed(x)for layer in self.transformer:x = layer(x)x = self.norm(x)x = self.head(x[:, 0])return x# 示例用法
if __name__ == "__main__":model = VisionTransformer()input_tensor = torch.randn(1, 3, 224, 224)  # 假设输入图像大小为224x224output = model(input_tensor)print(output.shape)  # 输出形状应为 (1, 1000)

五、模型优缺点

优点：

全局依赖关系：通过自注意力机制，ViT 能够捕获图像中的全局依赖关系，这对于理解复杂的视觉场景非常有用。
灵活的输入表示：ViT 的 patch-based 输入表示方式使得模型可以灵活地处理不同分辨率的图像。
强大的特征提取能力：Transformer 的强大建模能力使得 ViT 在大规模数据集上表现出色。
端到端训练：ViT 可以从头开始训练，不需要复杂的预处理步骤。

缺点：

计算成本高：ViT 的自注意力机制计算复杂度较高，特别是对于高分辨率图像，计算量和内存消耗都非常大。
数据需求大：ViT 需要在大规模数据集上进行训练才能取得良好的性能，对于小规模数据集的效果可能不如传统的卷积神经网络。
过拟合风险：由于模型参数量较大，ViT 在小规模数据集上容易发生过拟合。
训练不稳定：ViT 的训练过程可能不够稳定，需要仔细调整超参数和优化策略。

ViT模型相比于CNN架构模型优点在于它可以借助transformer全局信息互通、信息融合的特性来从全局的角度进行特征信息的提取，可以有效提高对复杂图像的理解能力，但图像信息却又不像文本信息那样有很强的上下文关联性，甚至图像缺少部分像素也不影响对图像的识别任务，所以这种全局的强关联性又是比较冗余的信息，同时还加大了运算量；即我们既希望模型可以多关注一点全局信息的特征，但又不希望过多的去关注全局的特征，这个是ViT模型所存在的问题。