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

时间：2024-11-21 18:58:38

       视觉常用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 here
        self.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/16

import torch
import torch.nn as nn
import torch.nn.functional as F

class 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_size
        self.patch_size = patch_size
        self.num_patches = (img_size // patch_size) ** 2
        self.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 x

class 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 src

class 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)  # 假设输入图像大小为224x224
    output = model(input_tensor)
    print(output.shape)  # 输出形状应为 (1, 1000)

五、模型优缺点

优点：

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

缺点：

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

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

相关文章