论文阅读（十三）：AN IMAGE IS WORTH 16X16 WORDS: TRANSFORMERS FOR IMAGE RECOGNITION AT SCALE

悟兰因w

已于 2024-10-08 20:50:16 修改

阅读量923

点赞数 9

分类专栏： # 论文阅读文章标签：论文阅读 python 深度学习神经网络网络人工智能

于 2024-10-07 16:16:39 首次发布

本文链接：https://blog.csdn.net/qq_51976293/article/details/142735136

版权

论文阅读专栏收录该内容

10 篇文章 0 订阅

订阅专栏

文章目录

摘要
1.ViT简介
2.ViT模型结构
3.ViT-B/16
4.Hybrid模型

论文：An image is worth 16x16 words: Transformers for image recognition at scale
论文链接：An image is worth 16x16 words: Transformers for image recognition at scale
论文代码：Github

摘要

Transformer架构在计算机视觉上的应用较为有限，在视觉方面，注意力机制要么用于和卷积神经网络相结合，要么用于替换卷积神经网络的某些组件。事实上，这种对CNN架构的依赖是不必要的，直接将图像块序列（sequences of image patches)）应用在Transformer模型上就能达到很好的图像分类效果。这种Transformer结构也称为Vision Transformer（视觉Transformer）。

1.ViT简介

受 NLP 中 Transformer 成功放缩/扩展 (scaling / scale up) 的启发，本文尝试将标准Transformer直接应用于图像，。为此，将图像拆分为块 (patch)，并将这些图像块的线性嵌入序列作为 Transformer 的输入。图像块 (patches) 的处理方式同 NLP 的令牌(tokens) (故经过线性嵌入后又叫 patch token)，以有监督方式训练图像分类模型。

token的含义：
NLP领域有token、embedding、encoding三大概念。

token：模型输入基本单元，可以是一个字、一个单词、一个词组、一个标点符号、一个字符等，取决于文本处理的需求和方法。在输入到模型前，文本通常需要转换成数值的形式才能作为输入传递给模型进行处理。故而，“token” 通常指的是将文本中的每个单词或其他单位转换为对应的数字或向量表示的过程。例如，"I love Java"就可被划分为三个token：“I”、“love”、“Java”，这些token可使用词嵌入（word embedding）或one-hot编码等方法转换为数值的形式。
tokenization：将文本划分为若干个token是文本处理的第一步，这个过程被称为 “tokenization”。

embedding（词嵌入）：一种将文本单位映射到连续向量空间的方法，文本单位被转换为实数向量后便于被计算机处理。具有相似语义的文本单位在向量空间中的距离较近，而语义不相似的单词在向量空间中的距离较远。这种连续向量表示可以捕捉到单词之间的语义和语法关系，从而在文本处理任务中能够更好地表示词语之间的相似性和差异性。常见的词嵌入方法如Word2Vec，其通过基于上下文的方法来学习单词的向量表示。
encoding（编码）：将输入数据转换为低维度、紧凑表示的过程。这种编码通常用于降维、特征提取、特征表示等任务，旨在从高维度的输入数据中提取有用的特征，并将其转换为更简洁、更可表达的形式，以便用于后续的机器学习、模型训练等任务。

在中等规模数据集（如ImageNet）上训练时，Vision Transformer (ViT)的精度可能比同规模的ResaNet低，这是因为Transformer缺乏CNN固有的一些归纳偏差 (inductive biases) ，如平移等效性和局限性（translation equivariance and locality），因此在数据量不足时，模型不能很好地泛化。
在更大规模数据集上训练时， ViT在大规模数据集上进行预训练，并迁移到具有较少数据点的任务上获得了出色结果。当在公共 ImageNet-21k 数据集或内部 JFT-300M 数据集上进行预训练时，ViT 在多个图像识别基准上接近或击败了最先进的技术。特别是，最佳模型在 ImageNet 上的准确率达到 88.55%，在ImageNet-RealL 上达到 90.72%，在 CIFAR-100 上达到94.55%，在 19 个任务的 VTAB 上达到77.63%。

2.ViT模型结构

补充知识：Imaging Patching

图像分块（Imaging Patching），指将图像划分为一系列大小相同或不同的小块，这些小块也称为图像块（Patch）。
在这里插入图片描述
将图像分为小块（Patch）带来的优势：

1.自适应性：一些自适应处理算法中可对不同图像区域采取不同的处理策略，将图像分块可使算法在局部区域上更加灵活。
2.特征提取：在一些任务中，特定区域的信息比整张图像更加有用，对每个Patch进行特征提取可获得更细粒度的信息，有助于更好理解图像内容。

补充知识：Patch Embedding

在这里插入图片描述
传统的卷积神经网络在图像处理中使用的是像素级操作，即通过卷积核在图像上滑动以提取特征。而图像块嵌入（Patch Embedding），将输入的图像分成小块（patch），之后将小块用向量序列表示，这些向量可作为模型的输入（这些小块本身也可看做是向量，完成特征提取后甚至是特征向量）。
图像块嵌入（Patch Embedding）的目的在于降低计算复杂度并提高特征提取的效率，因为传统卷积操作中，相邻像素常常会有大量重叠，而使用图像块嵌入将图像分块后可减少冗余计算并保留重要的特征信息。

补充知识：Class Token

类别令牌（Class Token），用于表示整个图像的类别信息，其通常会被添加到Patch Embedding得到的向量序列的某个位置，使得模型能够利用这一类别信息进行分类或生成任务。
【Add Class Token】
在Transformer模型中，Class Token一般会加在由Patch Embedding所得向量序列的开头，用于表示整个图像的类别信息，使得模型能对类别信息进行编码和利用。例如：
在这里插入图片描述
举一个更具体的例子，在一个图像二分类任务（是与不是）中，使用one-hot编码表示类别信息，有：

是：用向量[1，0]表示。
不是：用向量[0，1]表示。

假设Patch Embedding得到长度为196的向量序列，将这个Class Token与该序列连接得到最终的输入序列：

[Class_token, v1, v2, v3, ..., v196]

此时输入序列就包含了整张图象的类别信息，模型在训练过程中即可利用此类别信息帮助完成图像分类任务。

【Positional Encoding】
在NLP任务中，输入的是文本序列，为保留文本的位置信息，常添加位置编码（PE，Positional Encoding）。类似地，在Vision Transformer（ViT）中输入的是图像的Patch Embedding序列，位置编码则用于将Patch Embedding得到的向量序列中每个向量（对应不同的图像块）与其位置信息相关联，从而将整张图像的全局位置信息引入到Transformer模型中。
Vision Transformer（ViT）中常用PE(pos,2i)和PE(pos,2i+1)公式来计算位置编码：
在这里插入图片描述

在这里插入图片描述

其中，PE(pos,2i)对应维度为偶数的位置编码，PE(pos,2i+1)对应维度为奇数的位置编码。pos表示patch在序列中的位置，i是位置编码的维度索引（从0开始），dmodel是Transformer模型中的隐藏层维度（也称为特征维度）。
举例，假设有一张图像，其被分为4x4个图像块（patch），每个小块用一个2维向量表示，并设隐藏层维度（dmodel）为4，计算Class token和每个patch的位置编码。

d_model = 4
i = 0

PE(pos=0, 2i) = sin(0 / 10000^(2*0 / 4)) = sin(0) = 0
PE(pos=0, 2i + 1) = cos(0 / 10000^(2*0 / 4)) = cos(0) = 1

故Class token位置编码为[0，1]。
计算每个patch的位置编码（位置编号1-16）：

d_model = 4
i = 0, 1, 2, 3

pos = 1
PE(pos=1, 2*0) = sin(1 / 10000^(2*0 / 4)) = sin(1) ≈ 0.8415
PE(pos=1, 2*0 + 1) = cos(1 / 10000^(2*0 / 4)) = cos(1) ≈ 0.5403

pos = 2
PE(pos=2, 2*0) = sin(2 / 10000^(2*0 / 4)) = sin(2) ≈ 0.9093
PE(pos=2, 2*0 + 1) = cos(2 / 10000^(2*0 / 4)) = cos(2) ≈ -0.4161

以此类推，最终得到每个小块的位置编码结果。

2.1总体架构

在这里插入图片描述
在模型设计方面尽可能遵循原始Transformer架构，整体架构可分为三部分：

1.Linear Projection of Flattened Patches（Embedding层）：输入图像，经过Patch Embedding操作将其划分并展平为16x16的图像patches序列（Flattened Patches），并输入到Linear Projection of Flattened Patches层当中进行线性映射。每个patches都会得到对应的token，并在token序列首部加上class token来表示图像的类别信息（称为类别嵌入，Class Embedding，帮助模型获取图像整体信息）。之后token都会加上对应的位置编码（Position Embedding，形状相同的向量作为位置信息，与原token向量对应维度数据直接相加，图像中表示为0、1、2、3…，也称为位置嵌入）。

在这里插入图片描述

2.Transformer Encoder：将所有token及其位置参数输入到Transformer Encoder层中，该层是将上图部分堆叠了L次所得：

在这里插入图片描述

3.MLP Head（最终用于分类的层结构）：（多头注意力机制中有几个输入，就有几个输出）由于仅仅用于图像分类，故只将class token对应的输出提取并输入到MLP Head中得到分类的结果。

下文均采用ViT-B/16版本。

2.2Embedding层

2.2.1模型结构

在这里插入图片描述

对于标准的Transformer模块，其输入要求是token序列（向量序列），该序列可看作是大小为 $(t o k e n 个数, t o k e n 维数)$ 的二维矩阵。在代码实现中（以ViT-B/16为例），可直接使用卷积核大小为16x16、步长为16、卷积核个数为768（对应token向量的长度）的卷积层实现。尺度变化为：
$[224, 224, 3] - > [14, 14, 768] - > [196, 768]$
即，共196个Token，维度为768。
而在输入到Transformer Encoder层之前，还需加上class token（类别嵌入Class Embedding，帮助模型获取图像整体信息）以及位置编码，其中，位置编码是可训练的参数（位置嵌入，参数可在训练过程中调整，使得模型可学习到图像块patch在原始图像中的位置信息）。在拼接class token（尺度为 $[1, 768]$ ）后，尺度变化为：
$C a t ([1, 768], [196, 768]) - > [197, 768]$
再叠加位置编码（对应元素逐位相加）：
$[197, 768] - > [197, 768]$
尺度不发生变化（注意，拼接≠叠加）。
Transformer需要嵌入位置编码来保留输入图像patches之间的空间位置信息，论文中对于位置编码的使用进行了实验，以下分别为不使用、使用一维位置编码、使用二维位置编码、使用相对位置编码的实验结果：
在这里插入图片描述
可见，若不提供位置编码效果会差，但其它各种类型的编码效果效果都接近，这主要是因为 ViT 的输入是相对较大的图像块而非像素，所以学习位置信息相对容易很多。
除此之外还求了位置编码之间的余弦相似度来表明任意两个patches之间在位置上的关联程度（右侧以颜色的深浅表示相似程度）：
在这里插入图片描述
注意，论文中原图像大小为224x224，patches大小为32x32，由于224/32=7，故可将原图像划分为7x7的图像patches矩阵。比如，第一行第一列的小方块即表示该patch与其他patches的余弦相似度颜色矩阵，查看该方块（仍是7x7的小矩阵）：
在这里插入图片描述
可见，其与自身、同行、同列patch的位置编码相似度较高。

2.2.2代码详解

class PatchEmbed(nn.Module)：将输入图像分割成固定大小的块，并将每个块映射为一个向量表示，输出的结果适合输入给Transformer模型。

class PatchEmbed(nn.Module):
    """
    2D Image to Patch Embedding
    """
    #img_size:输入图像的尺寸;patch_size:图像patch的尺寸;in_c:输入图像的通道数;embed_dim:嵌入尺寸,即每个图像patch映射到的向量维度;norm_layer:可选的归一化层,若为None则不进行归一化
    def __init__(self, img_size=224, patch_size=16, in_c=3, embed_dim=768, norm_layer=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]
        #使用核大小16x16、步长为16的二维卷积操作相等于将图像分割为196个16x16的小块,再对每个小块内的像素进行卷积.由于行、列方向均能划分出14个小块用于卷积,故卷积操作后输出维度为(B,768,14,14)
        self.proj = nn.Conv2d(in_c, embed_dim, kernel_size=patch_size, stride=patch_size)
        #根据norm_layer初始图像归一化层(可能不进行归一化)
        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()

    def forward(self, x):
    	#获取批次大小、通道数、高度和宽度
        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]
        #self.proj对图像卷积后得到(B,768,14,14)的数据,通过flatten(2)得到形状为[B,768,14*14]的数据(相当于将图像卷积后得到的二维矩阵展平为了一维行向量序列)
        #通过transpose(1, 2)调整维度顺序,得到最终输出形状为[B,196,768](相当于得到一维列向量序列)
        x = self.proj(x).flatten(2).transpose(1, 2)
        x = self.norm(x)
        return x

class VisionTransformer(nn.Module)类下的forward_features(self, x)方法：将输入图像处理为特征向量序列，并完成类别嵌入（class embedding）、位置嵌入（position embedding）等过程。

    def forward_features(self, x):
        #将图像映射为向量序列,(B,3,224,224)->(B,196,768)
        x = self.patch_embed(x)  # [B, 196, 768]
        #将类别标记cls_token的形状由[1, 1, 768]扩展为[B, 1, 768]
        #判断是否使用蒸馏标记(并未使用)
        cls_token = self.cls_token.expand(x.shape[0], -1, -1)
        if self.dist_token is None:
        	#将类别标记与向量序列嵌入连接,[B,196,768]->[B,197,768]
            x = torch.cat((cls_token, x), dim=1)
        else:
            x = torch.cat((cls_token, self.dist_token.expand(x.shape[0], -1, -1), x), dim=1)
		#完成位置嵌入,并送入self.pos_drop()进行随机失活,防止过拟合
        x = self.pos_drop(x + self.pos_embed)
        #送入Encoder Block进行自注意力计算和特征提取
        x = self.blocks(x)
        #归一化处理
        x = self.norm(x)
        if self.dist_token is None:
        	#没有蒸馏标记时返回类别标记对应的输出,并经过self.pre_logits()得到最终分类结果
            return self.pre_logits(x[:, 0])
        else:
            return x[:, 0], x[:, 1]

class VisionTransformer(nn.Module)类下的forward(self, x)方法：调用forward(x)提取特征，并根据模型是否使用蒸馏标记决定如何生成最终输出。

def forward(self, x):
	#得到特征x
	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传入主分类头进行分类
		x = self.head(x)
	return x

其中，self.head = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity()是用于处理分类结果的全连接层（实际位于MLP中）。

2.3Transformer Encoder层

Transformer Encoder由多头注意力机制（Multi-Head Attention）和多层感知机模块（MLP Block）堆叠而成：
在这里插入图片描述
堆叠次数为L。且在每个块之前应用归一化层（Layer Norm层），每个块后应用Dropout层、残差连接 ( $⨁$ ，Residual Connection)。

Layer Norm：针对NLP领域提出的Normalization，此处是对每个token进行Norm处理。

MLP Block结构如下：
在这里插入图片描述
实际是全连接Linear+GELU激活函数+Dropout+全连接Linear+Dropout，经过第一个全连接层时参数变为原来的四倍，而第二个全连接层会还原参数个数。

2.3.2代码详解

Transformer Encoder层由Encoder Block模块堆叠而成，在class VisionTransformer(nn.Module)中表现为（depth=12）：

        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)
        ])

其中，class Block(nn.Module)定义为：

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)

        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

MLP模块的实现如下：

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.):
        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

除此之外，代码实现过程中还使用了DropPath（Stochastic Depth）来代替传统的Dropout结构。实现代码如下：

def drop_path(x, drop_prob: float = 0., training: bool = False):
    if drop_prob == 0. or not training:
        return x
    keep_prob = 1 - drop_prob
    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
    random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
    random_tensor.floor_()  # binarize
    output = x.div(keep_prob) * random_tensor
    return output


class DropPath(nn.Module):
    """
    Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
    """
    def __init__(self, drop_prob=None):
        super(DropPath, self).__init__()
        self.drop_prob = drop_prob

    def forward(self, x):
        return drop_path(x, self.drop_prob, self.training)

在代码中的调用方式如下：

self.drop_path = DropPath(drop_prob) if drop_prob > 0. else nn.Identity()

x = x + self.drop_path(self.token_mixer(self.norm1(x)))
x = x + self.drop_path(self.mlp(self.norm2(x)))

DropPath作为一种正则化手段，其效果是将深度学习模型中的多分支结构随机删除。例如，设置drop_out=0.2：

import torch

drop_prob = 0.2
keep_prob = 1 - drop_prob
#模拟输入数据
x = torch.randn(4, 3, 2, 2)
#模拟drop_path
shape = (x.shape[0],) + (1,) * (x.ndim - 1)
random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
random_tensor.floor_()
output = x.div(keep_prob) * random_tensor

得到x为：

tensor([[[[ 1.3833, -0.3703],
          [-0.4608,  0.6955]],
         [[ 0.8306,  0.6882],
          [ 2.2375,  1.6158]],
         [[-0.7108,  1.0498],
          [ 0.6783,  1.5673]]],

        [[[-0.0258, -1.7539],
          [-2.0789, -0.9648]],
         [[ 0.8598,  0.9351],
          [-0.3405,  0.0070]],
         [[ 0.3069, -1.5878],
          [-1.1333, -0.5932]]],

        [[[ 1.0379,  0.6277],
          [ 0.0153, -0.4764]],
         [[ 1.0115, -0.0271],
          [ 1.6610, -0.2410]],
         [[ 0.0681, -2.0821],
          [ 0.6137,  0.1157]]],

        [[[ 0.5350, -2.8424],
          [ 0.6648, -1.6652]],
         [[ 0.0122,  0.3389],
          [-1.1071, -0.6179]],
         [[-0.1843, -1.3026],
          [-0.3247,  0.3710]]]])

random_tensor为：

tensor([[[[0.]]],
        [[[1.]]],
        [[[1.]]],
        [[[1.]]]])

得到输出结果：

tensor([[[[ 0.0000, -0.0000],
          [-0.0000,  0.0000]],
         [[ 0.0000,  0.0000],
          [ 0.0000,  0.0000]],
         [[-0.0000,  0.0000],
          [ 0.0000,  0.0000]]],

        [[[-0.0322, -2.1924],
          [-2.5986, -1.2060]],
         [[ 1.0748,  1.1689],
          [-0.4256,  0.0088]],
         [[ 0.3836, -1.9848],
          [-1.4166, -0.7415]]],

        [[[ 1.2974,  0.7846],
          [ 0.0192, -0.5955]],
         [[ 1.2644, -0.0339],
          [ 2.0762, -0.3012]],
         [[ 0.0851, -2.6027],
          [ 0.7671,  0.1446]]],

        [[[ 0.6687, -3.5530],
          [ 0.8310, -2.0815]],
         [[ 0.0152,  0.4236],
          [-1.3839, -0.7723]],
         [[-0.2303, -1.6282],
          [-0.4059,  0.4638]]]])

dropout将随机drop_prob的数据置为0（失活），其余数据则进行缩放（乘以 $1/(1-drop\_prob)$ ）。dropout可用于提供网络的泛化能力、防止过拟合。

2.4MLP Head层

2.4.1模型结构

在这里插入图片描述
注意，在Transformer Encoder层前实际有个Dropout层，后有一个Layer Norm层，这些在源码中存在，但在论文的结构图中并未给出。 MLP Head用于获得最终的分类结果，由于只需要分类的信息，所以只需要提取出class token生成的对应结果即可。ViT-B/16中MLP Head实际结构如下：
在这里插入图片描述

2.4.2代码详解

在类class VisionTransformer(nn.Module)中实现了MLP-Head模块：

    def forward_features(self, x):
        # [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):
        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():
                # during inference, return the average of both classifier predictions
                return x, x_dist
            else:
                return (x + x_dist) / 2
        else:
            x = self.head(x)
        return x

def forward_features(self, x)：若不使用蒸馏标记，则执行return self.pre_logits(x[:, 0])，即输入到MLP中。
def forward(self, x)：若不适用蒸馏标记，则执行x = self.head(x)，而self.head = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity()实际是位于MLP中用于处理分类结果的全连接层。

而对于self.pre_logits()，其实现为：

self.pre_logits = nn.Sequential(OrderedDict([
                ("fc", nn.Linear(embed_dim, representation_size)),
                ("act", nn.Tanh())
            ]))

即为全连接Linear+激活函数Tanh。

3.ViT-B/16

以下是根据ViT-B/16源码绘制的模型结构图：
在这里插入图片描述

1.输入尺寸为 $(224, 224, 3)$ 的图像，经过Patch Embedding模块转换为向量序列。
- Conv2d：核大小16x16、步长16、共768个卷积核的卷积操作， $(224, 224, 3) - > (14, 14, 768)$ 。
- Fatten：在高度、宽度方向进行展平处理， $(14, 14, 768) - > (196, 768)$ 。
- Class token：将类别token与向量序列进行拼接， $(196, 768) - > (197, 768)$ 。
- Position Embedding：为所有token添加位置编码（也是可训练参数）， $(197, 768) - > (197, 768)$ （各向量对应维度元素相加，不改变数据维度）。
- Dropout层：随机关闭神经元，防止过拟合。
2.输入Transformer Encoder模块，由Encoder Block模块重复12次所得。
3.输入到Layer Norm层， $(197, 768) - > (197, 768)$ 。
4.Extract Class Token，即提取类别token的输出， $(197, 768) - > (1, 768)$ 。
5.MLP Head：若是在ImageNet21K数据集上训练，则Pre-Logits实际由Linear+tanh激活函数组成，而若在ImageNet1K或自己数据集上训练，则MLP Head只包含一个Linear即可。

在论文中共给出了ViT三个版本的模型参数：
在这里插入图片描述

Layers：Transformer Encoder中堆叠Encoder Block的次数。
Hidden Size：通过Embedding层后每个token的dim（向量长度）。
MLP size：Transformer Encoder中MLP Block第一个全连接的节点个数（Hidden Size的四倍）。
Heads：Transformer中Multi-Head Attention的heads数目。

4.Hybrid模型

Hybrid（混合）模型是指将传统CNN特征提取和Transformer进行结合。下图是将ResNet50作为特征提取器与ViT-B/16结合得到的混合模型：
在这里插入图片描述

R50 BackBone：进行特征提取。
- StdConv2d：使用核大小为7x7、步长为2、卷积核个数64的StdConv2d进行特征提取， $(224, 224, 3) - > (112, 112, 64)$ 。
- GroupNorm、ReLU、MaxPool：依次经过归一化、激活函数、最大池化完成特征提取， $(112, 112, 64) - > (56, 56, 64)$ 。
Stage1 Blockx3、Stage2 Blockx4、Stage3 Blockx9：在原Resnet50网络中，stage1重复堆叠3次，stage2重复堆叠4次，stage3重复堆叠6次，stage4重复堆叠3次，而此处的R50中，把stage4中的3个Block移至stage3中，所以stage3中共重复堆叠9次。

由于特征图大小为14x14x1024，无需再下采样，故Patch Embedding模块中使用大小为1x1、步长为1的卷积层，将特征矩阵的通道调整为768。之后过程与ViT-B/16模型相同。下表是论文用来对比ViT，Resnet（使用的卷积层和Norm层都进行了修改）以及Hybrid模型的效果。通过对比发现，在训练epoch较少时Hybrid优于ViT，但当epoch增大后ViT优于Hybrid：
在这里插入图片描述