Vision in Transformer

原理

image-20240318161334702

从架构图可以看出，ViT主要步骤如下：

  1. Patch Embedding：首先，对原始输入图像作切块处理。假设输入的图像大小为224×224，我们将图像切成一个个固定大小为16×16的方块，每一个小方块就是一个patch，那么每张图像中patch的个数为(224×224)/(16×16) = 196个。切块后，我们得到了196个[16, 16, 3]的patch，我们把这些patch送入Linear Projection of Flattened Patches(Embedding层)，这个层的作用是将输入序列展平。所以输出后也有196个token，每个token的维度经过展平后为16×16×3 = 768，所以输出的维度为[196, 768]。不难看出，Patch Embedding的作用是将一个CV问题通过切块和展平转化为一个NLP问题。
  2. Position Embedding: 我们知道，图像的每个patch和文本一样，也有先后顺序，是不能随意打乱的，所以我们需要再给每个token添加位置信息。类比BERT模型，我们还需要添加一个特殊字符class token。那么，最终要输入到Transformer Encoder的序列维度为[197, 768]。Position Embedding的作用是添加位置信息。
  3. Transformer Encoder:将维度为[197, 768]的序列输入到标准的Transformer Encoder中。
  4. MLP Head：Transformer Encoder的输出其实也是一个序列，但是在ViT模型中只使用了class token的输出，将其送入MLP模块中，最终输出分类结果。MLP Head的作用是用于最终的分类。

动图版本：

vit

程序流程图：

在这里插入图片描述

代码分析

drop_path技术

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
def drop_path(x, drop_prob: float = 0., training: bool = False):
    """
    对每个样本应用Drop paths（随机深度），通常用于残差块的主路径上。


    将层和参数名称更改为'drop path'，而不是将DropConnect作为层名称并使用
    '生存率'作为参数。
    
    参数：
    x (Tensor): 输入的张量。
    drop_prob (float): 丢弃概率，定义了每个元素被丢弃的概率。
    training (bool): 指示模型是否处于训练模式。

    返回值：
    Tensor: 经过drop path处理的张量。
    """
    if drop_prob == 0. or not training:
        # 如果不进行drop或不处于训练模式，则直接返回原张量
        return x
    keep_prob = 1 - drop_prob  # 计算保留概率
    # 生成与x的第一个维度相同的形状，其他维度为1，以适配不同维度的张量，不仅仅是2D的卷积网络
    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的值，并应用生成的mask
    return output


class DropPath(nn.Module):
    """
    对每个样本应用Drop paths（随机深度），通常用于残差块的主路径上。
    这个类是一个nn.Module，它可以作为神经网络中的一部分来使用，以便在训练过程中应用随机深度策略。
    
    参数：
    drop_prob (float, 可选): 丢弃概率，定义了每个元素被丢弃的概率。如果为None，则不应用drop path。
    """
    def __init__(self, drop_prob=None):
        """
        初始化DropPath模块。
        
        参数：
        drop_prob (float, 可选): 丢弃概率，定义了每个元素被丢弃的概率。如果为None，则此模块不会执行任何操作。
        """
        super(DropPath, self).__init__()
        self.drop_prob = drop_prob

    def forward(self, x):
        """
        定义模块的前向传播。
        
        参数：
        x (Tensor): 输入的张量。
        
        返回值：
        Tensor: 经过drop path处理的张量，或在不应用drop path时返回原始输入张量。
        """
        # 如果没有设置丢弃概率，则直接返回输入的x，相当于这个模块不进行任何操作
        # 否则，调用drop_path函数应用随机深度策略
        return drop_path(x, self.drop_prob, self.training)

PatchEmbed层

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
class PatchEmbed(nn.Module):
    """
    将2D图像转换为批量嵌入。
    
    参数：
    img_size (int): 输入图像的大小，默认为224。假设输入图像为正方形。
    patch_size (int): 每个补丁的大小，默认为16。补丁也是正方形的。
    in_c (int): 输入图像的通道数，默认为3（对应于RGB图像）。
    embed_dim (int): 嵌入向量的维度。
    norm_layer (nn.Module): 对嵌入向量进行归一化的层。如果为None，则使用nn.Identity()。
    """
    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]

        # 使用卷积层来实现图像到补丁的投影，卷积核大小和步长均为补丁大小
        # in_c=3,embed_dim=768,kernel_size和stride都是16
        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):
        # 提取输入x的形状：批量大小B，通道数C(3)，高度H(14)，宽度W(14)
        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]})."

        # 使用定义好的卷积层proj对图像进行处理，然后将结果平铺并转置，以便于后续处理
        # 经过proj(x)后
        x = self.proj(x).flatten(2).transpose(1, 2)
        # 对嵌入向量进行归一化处理
        x = self.norm(x)
        return x

重点是分析proj这个卷积层的作用：

输入

  • 输入数据: 输入到这个卷积层的数据是一个四维张量，形状为[B, C, H, W]，其中：
    • B是批量大小，表示一次处理多少个图像。
    • C是通道数，对于彩色RGB图像，C=3。
    • H和W是图像的高度和宽度,都是244

卷积层配置

  • 卷积核大小 (kernel_size): 与补丁的大小相同，即16x16。这意味着每次卷积操作都会覆盖16x16的像素区域。
  • 步长 (stride): 同样设置为补丁的大小，即16。这确保了每次卷积操作之后都会移动一个补丁的大小，没有重叠部分，从而直接将图像分割成大小相等的补丁。
  • 输入通道数 (in_channels): 等于输入数据的通道数，对于RGB图像来说是3。
  • 输出通道数 (out_channels): 这是嵌入维度，例子中是768。这表示每个补丁都会被映射到一个768维的向量。

输出

  • 输出数据: 经过卷积层处理后，输出数据的形状将会是[B, embed_dim, grid_H, grid_W]，其中：
    • B是批量大小，与输入相同。
    • embed_dim是输出通道数，即嵌入维度，你的例子中是768。
    • grid_H和grid_W是经过卷积层处理后的图像高度和宽度在补丁维度上的大小。由于原始图像尺寸为224x224，补丁大小为16x16，所以grid_H和grid_W都是14（224/16=14），这意味着每个维度上都有14个补丁

展平并转置

  • flatten(2)操作是将每个补丁的嵌入向量展平。这里的2指的是从第三维开始平铺（因为维度是从0开始计数的，所以0对应于批次维度B，1对应于嵌入维度embed_dim）。因此，flatten(2)操作是将每个补丁的嵌入向量平铺开来，结果的形状变为[B, embed_dim, num_patches]，其中num_patches是总的补丁数量（即grid_size[0] * grid_size[1]）
  • 最后，.transpose(1, 2)操作是将第二维和第三维交换。由于在flatten(2)之后，我们得到的维度是[B, embed_dim, num_patches]，我们需要将嵌入向量的维度（embed_dim）和补丁数量的维度（num_patches）交换，以使得每个补丁的嵌入向量在序列的第二维度上。这样做的目的是为了让输出的形状与后续Transformer网络的输入要求相匹配。转置后，输出的形状是[B, num_patches, embed_dim]。

Decoder-Attention层

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
class Attention(nn.Module):
    def __init__(self,
                 dim,  # 输入token的维度
                 num_heads=8,  # 注意力头的数量
                 qkv_bias=False,  # 是否在QKV线性变换中加入偏置
                 qk_scale=None,  # QK缩放因子，如果为None，则默认为head_dim ** -0.5
                 attn_drop_ratio=0.,  # 注意力权重的dropout比率
                 proj_drop_ratio=0.):  # 输出投影的dropout比率
        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)
        # 注意力权重的dropout
        self.attn_drop = nn.Dropout(attn_drop_ratio)
        # 定义输出的线性变换FFN
        self.proj = nn.Linear(dim, dim)
        # 输出投影的dropout
        self.proj_drop = nn.Dropout(proj_drop_ratio)

    def forward(self, x):
        # 输入形状 [batch_size, num_patches + 1, total_embed_dim]
        B, N, C = x.shape

        # 通过qkv线性变换并重塑和置换得到Q, K, V
        
        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
        
        # 输出形状: [3, batch_size, num_heads, num_patches + 1, embed_dim_per_head]
        
        # 分解为Q, K, V,形状为[batch_size, num_heads, num_patches+1, embed_dim_per_head]
        q, k, v = qkv[0], qkv[1], qkv[2]
         
        # 计算注意力分数，应用缩放，执行Softmax得到注意力权重
        # 交换了最后两个元素的位置后相乘：[batch_size,num_heads,num_patches+1,num_patches+1]
        attn = (q @ k.transpose(-2, -1)) * self.scale
        attn = attn.softmax(dim=-1)
        attn = self.attn_drop(attn)  # 应用注意力dropout

        # 通过注意力权重和V计算输出，然后进行重塑和线性变换
        # attn@v的结果是[batch_size,num_heads,num_patches+1,embed_dim_per_head]
        # 交换了第2,3个元素位置：[batch_size,num_patches+1,num_heads,embed_dim_per_head]
        # 最后变换为：[batch_size,num_patches+1,num_heads*embed_dim_per_head]
        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
        x = self.proj(x)
        x = self.proj_drop(x)  # 应用输出投影的dropout
        return x

Decoder-MLP层

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
class Mlp(nn.Module):
    """
    在Vision Transformer、MLP-Mixer及相关网络中使用的多层感知机（MLP）。
    
    参数:
    in_features (int): 输入特征的数量。
    hidden_features (int, 可选): 隐藏层特征的数量。如果未指定，默认使用与输入特征相同的数量。
    out_features (int, 可选): 输出特征的数量。如果未指定，默认使用与输入特征相同的数量。
    act_layer (nn.Module, 可选): 激活函数层，默认为GELU。
    drop (float, 可选): Dropout比率，默认为0（不应用dropout）。
    """
    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)
        # Dropout层
        self.drop = nn.Dropout(drop)

    def forward(self, x):
        # 应用第一层线性变换
        x = self.fc1(x)
        # 应用激活函数
        x = self.act(x)
        # 应用dropout
        x = self.drop(x)
        # 应用第二层线性变换
        x = self.fc2(x)
        # 再次应用dropout
        x = self.drop(x)
        return x

Decoder层

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
class Block(nn.Module):
    """
    Vision Transformer中的一个基本块，包括自注意力层和MLP层。
    
    参数:
    dim (int): 特征维度。
    num_heads (int): 自注意力机制中头的数量。
    mlp_ratio (float): MLP隐藏层大小相对于输入层的比例，默认为4。
    qkv_bias (bool): 是否在QKV计算中添加偏置。
    qk_scale (float): 自注意力中QK缩放因子，如果为None，则默认为dim的平方根的倒数。
    drop_ratio (float): 在MLP和自注意力投影后应用的dropout比率。
    attn_drop_ratio (float): 在计算自注意力权重后应用的dropout比率。
    drop_path_ratio (float): 应用于随机深度（DropPath）的比率。
    act_layer (nn.Module): 激活函数，默认为GELU。
    norm_layer (nn.Module): 归一化层，默认为LayerNorm。
    """
    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)
        # 随机深度DropPath
        self.drop_path = DropPath(drop_path_ratio) if drop_path_ratio > 0. else nn.Identity()
        self.norm2 = norm_layer(dim)  # 第二层归一化
        # MLP层，其隐藏层大小是输入特征维度的mlp_ratio倍
        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):
        # 应用自注意力层，包括归一化、自注意力、DropPath
        x = x + self.drop_path(self.attn(self.norm1(x)))
        # 应用MLP层，包括归一化、MLP、DropPath
        x = x + self.drop_path(self.mlp(self.norm2(x)))
        return x

权重初始化

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
def _init_vit_weights(m):
    """
    对Vision Transformer模型中的权重进行初始化。

    参数:
    m (nn.Module): 需要初始化权重的模块。

    方法:
    - 对于线性层（nn.Linear），使用截断正态分布（truncated normal distribution）初始化权重，标准差为0.01。如果线性层包含偏置，则将偏置初始化为0。
    - 对于卷积层（nn.Conv2d），使用Kaiming正态初始化方法（也称为He初始化），模式为"fan_out"，这适用于ReLU激活函数之后的卷积层。如果卷积层包含偏置，则将偏置初始化为0。
    - 对于层归一化（nn.LayerNorm），将偏置初始化为0，权重初始化为1。
    """
    if isinstance(m, nn.Linear):
        # 使用截断正态分布初始化线性层的权重，标准差设置为0.01
        nn.init.trunc_normal_(m.weight, std=.01)
        # 如果线性层有偏置，则将偏置初始化为0
        if m.bias is not None:
            nn.init.zeros_(m.bias)
    elif isinstance(m, nn.Conv2d):
        # 使用Kaiming正态初始化卷积层的权重，适用于ReLU激活后的情况
        nn.init.kaiming_normal_(m.weight, mode="fan_out")
        # 如果卷积层有偏置，则将偏置初始化为0
        if m.bias is not None:
            nn.init.zeros_(m.bias)
    elif isinstance(m, nn.LayerNorm):
        # 将层归一化的偏置初始化为0
        nn.init.zeros_(m.bias)
        # 将层归一化的权重初始化为1
        nn.init.ones_(m.weight)

整体架构

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
class VisionTransformer(nn.Module):
    """
    Vision Transformer模型。
    
    参数:
    img_size (int, tuple): 输入图像的尺寸。
    patch_size (int, tuple): 图像分块的尺寸。
    in_c (int): 输入通道数。
    num_classes (int): 分类任务的类别数。
    embed_dim (int): 嵌入（特征）维度。
    depth (int): Transformer的层数。
    num_heads (int): 注意力机制的头数。
    mlp_ratio (float): MLP隐藏层维度与嵌入维度的比例。
    qkv_bias (bool): QKV计算中是否添加偏置项。
    qk_scale (float): 自注意力中的缩放因子。
    representation_size (Optional[int]): 如果设置，表示在分类头之前添加一个表示层。
    distilled (bool): 是否包括蒸馏token和头，如DeiT模型。
    drop_ratio (float): dropout比率。
    attn_drop_ratio (float): 注意力权重的dropout比率。
    drop_path_ratio (float): 随机深度（DropPath）的比率。
    embed_layer (nn.Module): 用于实现图像到patch嵌入的层。
    norm_layer (nn.Module): 归一化层。
    act_layer (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):
        super(VisionTransformer, self).__init__()
        self.num_classes = num_classes
        self.num_features = self.embed_dim = embed_dim  # 为了与其他模型保持一致
        self.num_tokens = 2 if distilled else 1  # 根据是否蒸馏，确定token的数量
        norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)
        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

        # 分类token和可选的蒸馏token
        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        self.dist_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) if distilled else None
        # 位置嵌入
        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim))
        self.pos_drop = nn.Dropout(p=drop_ratio)

        # Transformer块
        dpr = [x.item() for x in torch.linspace(0, drop_path_ratio, depth)]  # 计算drop path比率
        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
            self.pre_logits = nn.Sequential(OrderedDict([
                ("fc", nn.Linear(embed_dim, representation_size)),
                ("a"
                 "ct", 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()
        self.head_dist = None
        if distilled:
            self.head_dist = nn.Linear(self.embed_dim, self.num_classes) if num_classes > 0 else nn.Identity()

        # 权重初始化
        self._init_weights()

    def _init_weights(self):
        # 权重初始化的实现省略
        pass

    def forward_features(self, x):
        # 将图像通过patch嵌入层转换
        x = self.patch_embed(x)
        # 添加分类token和可选的蒸馏token
        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)
        else:
            x = torch.cat((cls_token, self.dist_token.expand(x.shape[0], -1, -1), x), dim=1)

        # 应用位置嵌入和dropout
        x = self.pos_drop(x + self.pos_embed)
        # 通过一系列Transformer块
        x = self.blocks(x)
        # 应用最后的归一化
        x = self.norm(x)
        # 根据是否有蒸馏token选择输出
        if self.dist_token is None:
            #这里，x[:, 0]实际上就是从Transformer的输出中选取clstoken对应的特征向量，形状为[B, C]
            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():
                return x, x_dist
            else:
                return (x + x_dist) / 2
        else:
            x = self.head(x)
        return x
  • 表示层（Representation Layer）：如果representation_size被设置，这表明在模型的最后，希望通过一个额外的全连接层（fc）和激活函数（act，这里使用的是Tanh）来处理logits。这可以用于降维或学习一个更加紧凑的特征表示。如果不进行蒸馏（distilled为False），这个层会被激活。
  • 分类头（Classification Head）：这是模型的最后一层，负责将特征表示（logits）映射到类别预测上。这里使用的是一个线性层nn.Linear，其输出大小为num_classes，表示不同类别的预测logits
  • 蒸馏头（Distillation Head）：如果进行模型蒸馏（distilled为True），则会添加一个蒸馏头head_dist。蒸馏头的作用是在训练时学习从教师模型转移过来的知识，通常用于提高模型的泛化能力或压缩模型。

分类的具体过程：

当模型的前向传播方法forward_features被调用时，它返回了一个包含cls_token和所有补丁嵌入的合并序列。这个序列的形状为[B, N+1, C]，其中：

  • B是批次大小（batch size），
  • N是补丁的数量，
  • C是嵌入的维度（embed_dim），
  • N+1中的+1代表了额外的cls_token。

在代码中，x[:, 0]用于选取每个序列中位置0上的元素，即cls_token对应的特征向量。由于cls_token被放在序列的第一个位置，x[:, 0]能够精确地选取出这个全局表示向量。这样做的结果是，对于批次中的每个图像，我们都得到了一个代表其全局信息的向量，形状为[B, C]：

  • B代表批次中图像的数量，
  • C代表特征向量的维度。

接下来，这个[B, C]形状的特征向量被送入分类头（即全连接层self.head），进行最终的类别预测。全连接层将特征向量从C维映射到类别数num_classes维，输出的形状变为[B, num_classes]，其中每一行包含了对应图像属于每个类别的预测分数。

总结

在Vision Transformer (ViT)中，预测过程是这样进行的：

输入图像预处理

  • 输入: 首先，输入图像的尺寸为224x224，颜色通道为3（RGB）。因此，每个图像的原始张量形状为[3, 224, 224]。

  • 批处理: 如果我们有N个这样的图像，组成一个批次进行处理，那么输入张量的形状将为[N, 3, 224, 224]。

图像切分为补丁

  • 补丁化: 每个图像被切分成16x16的小块，即补丁。因为每边可以分为14个这样的补丁（224 / 16 = 14），总共有14 * 14 = 196个补丁。

  • 补丁形状: 每个补丁被展平，展平后的维度为3 * 16 * 16 = 768，因此每个图像转换为196个这样的768维向量。

  • 补丁化后的形状: 对整个批次而言，补丁化后的张量形状为[N, 196, 768]。

类别嵌入

  • 在补丁序列的前面加入一个特殊的类别嵌入（类似于BERT中的[CLS]标记），这个嵌入是可学习的，旨在最终被用于分类。

  • 添加类别嵌入后: 张量的形状变为[N, 197, 768]，因为每个图像序列前都添加了一个额外的768维向量。

通过Transformer编码器

  • 编码: 这197个向量（包括196个补丁向量和1个类别嵌入向量）被送入一系列Transformer编码器层。这些层通过自注意力机制和前馈网络，能够捕捉补丁之间的复杂关系。

  • 保持形状不变: 经过Transformer层处理后，输出张量的形状仍为[N, 197, 768]，但是包含了丰富的上下文信息。

预测头

  • 提取类别嵌入: 对于分类任务，我们关注的是类别嵌入（即序列中的第一个向量）。这个向量现在包含了经过整个图像处理和编码后的全局信息。

  • 预测: 类别嵌入被送入一个预测头（通常是一个或多个全连接层），这个预测头将768维的向量映射到最终的类别数上，假设这里是5。

    • 输出形状: 经过预测头处理后，每个图像的输出形状为[5]，代表5个类别的概率分布。

    • 批次输出形状: 对整个批次而言，输出张量的形状为[N, 5]。

结果解析

  • 概率分布: 输出张量的每一行代表一个图像对应的5个类别的概率分布。这个分布可以通过softmax函数获得，确保所有概率之和为1。
  • 分类决策: 最终的分类决策通常是选取概率最高的类别作为预测结果

总结来说，从输入图像到最终的分类预测，Vision Transformer通过将图像切分成补丁、添加类别嵌入、利用Transformer编码器捕捉补丁间的关系，最后通过预测头输出每个类别的概率，从而实现从图像到类别标签的映射。