视觉Transformer详解

自2017年提出《Attention is All You Need》¹以来，transformer模型已经成为自然语言处理（NLP）领域的标准技术。2021年，《An Image is Worth 16x16 Words》²一文成功地将transformer应用到计算机视觉任务。自此之后，许多基于transformer的架构被提出，用于计算机视觉任务。

本文介绍了《一张图片等于16x16个单词》² 中所述的视觉Transformer模型（ViT）。本文不仅包含ViT的开源代码，还详细解释了各个组件的概念。所有代码均基于PyTorch Python包。

照片由 Sahand Babali 拍摄，来自 Unsplash

本文是一系列深入探讨视觉Transformer内部工作机制的文章中的一部分文章。每一篇文章也都作为包含可执行代码的 Jupyter Notebook 提供。该系列中的其他文章有：

• 视觉Transformer详解
  → Jupyter Notebook
• 视觉Transformer注意力详解
  → Jupyter Notebook
• 视觉Transformer位置编码详解
  → Jupyter Notebook
• Token-to-Token视觉Transformer详解
  → Jupyter Notebook
• 视觉Transformer详解系列的GitHub仓库页面

• 什么是视觉Transformer？
• 模型概览
  — 图像标记
  — 标记处理
  — 编码模块
  — 神经网络模块
  — 预测
• 完整代码
• 结论
  — 更多阅读
  — 参考文献

如《注意力机制就是你所需要的》¹ 所介绍，Transformer 是一种主要通过注意力机制进行学习的机器学习模型。Transformer 很快成为这些任务中的最新技术，如语言翻译等序列到序列类型的任务。

_一张图等于16x16个单词_²成功地修改了[1]中提出的变压器模型，以解决图像分类任务，创建了Vision Transformer（ViT）。ViT基于[1]中变压器相同的注意力机制。然而，尽管用于NLP任务的变压器由编码器注意力机制和解码器注意力机制组成，ViT仅使用编码器部分。编码器的输出传递给一个神经网络“头”进行预测。

ViT在[2]中的实现存在的一个问题是，其最佳性能需要在大规模数据集上进行预训练。在专有JFT-300M数据集上预训练的模型表现最好。在较小的开源ImageNet-21k数据集上预训练的模型表现与最先进的卷积ResNet模型相当。

_从头开始在ImageNet上训练视觉变换器的Tokens-to-Token ViT_³ 试图通过引入一种新颖的预处理方法来消除对预训练的需求，该方法将输入图像转换成一系列token。更多关于此方法的信息，可以在这里找到here。本文将专注于讨论在[2]中实现的ViT。

本文遵循了《一张图值十六个字》² 这篇文章所描述的模型结构。然而，该论文的代码并未公开发布。较新的《Token-to-Token ViT》³ 论文的代码可在 GitHub 上获取。Tokens-to-Token ViT (T2T-ViT) 模型在普通 ViT 模型基础上添加了一个 Tokens-to-Token (T2T) 模块。本文的代码基于 Tokens-to-Token ViT ³ __ GitHub 代码中的 ViT 组件。本文对代码进行了以下修改，包括但不限于支持任意形状的输入图像以及去掉了 dropout 层。

下面是一个关于ViT模型的图解。

ViT 模型示意图 (作者画的图片)

ViT的第一步是将输入图像分割成tokens。Transformer是在一个_序列_的_tokens_上工作的；在自然语言处理中，这通常是一个由_单词_组成的句子。对于计算机视觉而言，如何将输入分割成tokens还不太明确。

这个 __ViT 将图像分割成 token，每个 token 对应图像的一个局部区域（或称为 patch）。他们这样描述：将高度为 H，宽度为 W，通道数为 C 的图像重塑为 N 个 token，每个 token 的 patch 大小为 P：

每个标记的长度为 P²∗C。

让我们来看一下 Luis Zuno (@ansimuz) 的像素艺术作品《名为 Mountain at Dusk》的补丁分词化示例。原始艺术品已被裁剪并转换成了单通道图像。这意味着每个像素的值都在0到1的范围内。单通道图像通常以灰度显示；然而，我们将用紫色来显示，因为这样更易于观察。

补丁分词处理未包含在[3]中的代码内。本节中的所有代码均为作者原创编写。

    mountains = np.load(os.path.join(figure_path, 'mountains.npy'))  

    H = mountains.shape[0]  
    W = mountains.shape[1]  
    print('黄昏中的山脉图像尺寸为 H =', H, '和 W =', W, '像素。')  
    print()  

    fig = plt.figure(figsize=(10,6))  
    plt.imshow(mountains, cmap='Purples_r')  
    plt.xticks(np.arange(-0.5, W+1, 10), labels=np.arange(0, W+1, 10))  
    plt.yticks(np.arange(-0.5, H+1, 10), labels=np.arange(0, H+1, 10))  
    plt.clim([0,1])  
    cbar_ax = fig.add_axes([0.95, .11, 0.05, 0.77])  # 将颜色条轴添加到图像中
    plt.colorbar(cax=cbar_ax);  
    #plt.savefig(os.path.join(figure_path, 'mountains.png'))

    黄昏时的山高 60 像素 宽 100 像素。

代码生成的图片，作者供图

这张图的高度是 H=60，宽度是 W=100。我们将把 P 设为 20，因为它可以同时平分 H 和 W。

P = 20  
N = int((H*W)/(P**2))  
print('将会生成', N, '个补丁，每个补丁的大小为', P, '乘以', P, '.')  
print()  

fig = plt.figure(figsize=(10,6))  
plt.imshow(mountains, cmap='Purples_r')  
plt.hlines(np.arange(P, H, P)-0.5, -0.5, W-0.5, color='w')  
plt.vlines(np.arange(P, W, P)-0.5, -0.5, H-0.5, color='w')  
plt.xticks(np.arange(-0.5, W+1, 10), labels=np.arange(0, W+1, 10))  
plt.yticks(np.arange(-0.5, H+1, 10), labels=np.arange(0, H+1, 10))  
生成一系列从9.5到W，步长为P的数组，重复3次
生成一系列从9.5到H，步长为P的数组，重复5次
对于从1到N的每个i:  
plt.text(x_text[i-1], y_text[i-1], str(i), color='w', fontsize='xx-large', ha='center')  
plt.text(x_text[2], y_text[2], str(3), color='k', fontsize='xx-large', ha='center');  
# plt.savefig(os.path.join(figure_path, 'mountain_patches.png'), bbox_inches='tight')

会有15个20乘20的补丁包，每个都是20乘20的大小。

代码结果（由作者提供图片）

通过平铺这些区域，我们看到了生成的标记。我们以第12个区域为例，因为它包含了四种不同的色调。

    print('每个补丁都会生成一个长度为' + str(P**2) + '的token。')  
    print('\n')  

    patch12 = mountains[40:60, 20:40]  
    token12 = patch12.reshape(1, P**2)  

    fig = plt.figure(figsize=(10,1))  
    plt.imshow(token12, aspect=10, cmap='Purples_r')  
    plt.clim([0,1])  
    plt.xticks(np.arange(-0.5, 401, 50), labels=np.arange(0, 401, 50))  
    plt.yticks([]);  
    #plt.savefig(os.path.join(figure_path, 'mountain_token12.png'), bbox_inches='tight')

每个补丁都会创建一个长度为400的凭证。

代码生成的图片（作者提供的图片）

提取图像中的 token 后，通常会使用线性变换来改变 token 的长度。这通常通过一个可学习的线性层来实现。新的 token 长度被称为 潜在维度、通道维 或 token 长度。经过变换后，这些 token 就不再能直接看出是来自原始图像的补丁了。

既然我们理解了这个概念，我们来看看代码中是如何实现补丁分词的。

    class Patch_Tokenization(nn.Module):  
        def __init__(self,  
                    img_size: tuple[int, int, int]=(1, 1, 60, 100),  
                    patch_size: int=50,  
                    token_len: int=768):  

            """ Patch Tokenization Module  
                Args:  
                    img_size (元组[int, int, int])：输入尺寸（通道数，高度，宽度），即图像的尺寸  
                    patch_size (int)：补丁的边长（为方形）  
                    token_len (int)：令牌长度  
            """  
            super().__init__()  

            ## 定义层结构  
            self.img_size = img_size  
            C, H, W = self.img_size  
            self.patch_size = patch_size  
            self.token_len = token_len  
            assert H % self.patch_size == 0, '图像的高度必须能被补丁尺寸整除，否则会抛出异常.'  
            assert W % self.patch_size == 0, '图像的宽度必须能被补丁尺寸整除，否则会抛出异常.'  
            self.num_tokens = (H / self.patch_size) * (W / self.patch_size)  

            ## 定义层结构  
            self.split = nn.Unfold(kernel_size=self.patch_size, stride=self.patch_size, padding=0)  
            self.project = nn.Linear((self.patch_size**2)*C, token_len)  

        def forward(self, x):  
            x = self.split(x).transpose(1,0)  
            x = self.project(x)  
            return x

注意，这两个assert语句确保图像尺寸能被补丁大小整除。补丁的实际分割是通过torch.nn.Unfold层来实现的。

我们将使用裁剪过的单通道版本的《黄昏山景》⁴来运行这段代码。我们应该看到与之前相同的token数量和初始token大小的值。我们将使用_token_len=768_作为投影长度，这与ViT²基础变体的大小一致。

以下代码块中的第一行将 "《日落之山》" 的数据类型从 NumPy 数组更改为 Torch 张量。我们将使用 unsqueeze 函数对张量进行操作以创建通道维度和批次大小维度。正如前面提到的，我们只需要一个通道。由于只有一个图像，因此 batchsize=1。

    x = torch.from_numpy(mountains).unsqueeze(0).unsqueeze(0).to(torch.float32)  
    token_len = 768  
    print('输入维度如下：\n\t批大小:', x.shape[0], '\n\t输入通道数:', x.shape[1], '\n\t图像尺寸:', (x.shape[2], x.shape[3]))  

    # 定义模块如下：  
    patch_tokens = Patch_Tokenization(img_size=(x.shape[1], x.shape[2], x.shape[3]),  
                                        patch_size = P,  
                                        token_len = token_len)

输入维度如下：
每个批次的批量大小：1  
输入通道数：1个  
图像尺寸：(60, 100)

现在，我们将图像分割成若干标记/碎片。

    x = patch_tokens.split(x).transpose(2,1)  
    print('补丁标记化之后，维度如下：\n\t批次大小：', x.shape[0], '\n\t标记个数：', x.shape[1], '\n\t标记维度：', x.shape[2])

    在分词处理后，维度如下  
       批量大小：1   
       token数量：15   
       每个token的长度：400.

如我们在示例中所见，共有_N=15_个标记，每个长度为400。最后，将这些标记投影为长度_tokenlen。

    x = patch_tokens.project(x)  
    print('投影后，维度如下：\n\t批次大小：', x.shape[0], '\n\ttoken的数量：', x.shape[1], '\n\ttoken的长度：', x.shape[2])

    投影之后，维度是  
       批量大小: 1   
       标记数量: 15   
       token长度: 768

我们现在有了 token，就可以进入 ViT 的步骤了。

我们将ViT编码块之前的两个步骤指定为“令牌处理阶段”。下面展示了ViT图中的令牌处理部分。

ViT Token处理组件图(作者绘制)

第一步是向图像标记前添加一个空白标记，称为预测令牌。这个标记将在编码块输出时用于进行预测。它一开始是空白的（或等同于零），这样它就能从其他图像令牌中获取信息。

我们将从175个token开始。每个token的长度为768，这是ViT²基础版本的维度。我们选择了13作为批次大小，因为它是一个素数，不会被误认为是其他参数。

    # 定义输入
    num_tokens = 175
    token_len = 768
    batch = 13
    x = torch.rand(batch, num_tokens, token_len)
    打印('输入维度为\n\t批量大小:', x.shape[0], '\n\ttoken个数:', x.shape[1], '\n\ttoken长度:', x.shape[2])

    # 添加预测token
    pred_token = torch.zeros(1, 1, token_len).expand(-1, -1, -1)
    打印('预测token的维度为\n\t批量大小:', pred_token.shape[0], '\n\ttoken个数:', pred_token.shape[1], '\n\ttoken长度:', pred_token.shape[2])

    x = torch.cat((pred_token, x), dim=1)
    打印('加上预测token后的维度为\n\t批量大小:', x.shape[0], '\n\ttoken个数:', x.shape[1], '\n\ttoken长度:', x.shape[2])

    输入维度如下  
       批量大小: 13  
       令牌数: 175  
       令牌长度: 768  
    预测令牌的维度如下  
       批量大小: 13  
       令牌数: 1  
       令牌长度: 768  
    包含预测令牌的维度如下  
       批量大小: 13  
       令牌数: 176  
       令牌长度: 768

现在，我们为我们的 token 添加了一个位置编码。位置编码使 transformer 能够理解图像 token 的位置。需要注意的是，这是相加，而不是相连接。关于位置编码的细节，留待另一个时间解释。

    def get_sinusoid_encoding(num_tokens, token_len):  
        """ 生成正弦位置编码表  

            参数:  
                num_tokens (int): token的数量  
                token_len (int): token的长度  

            返回:  
                (torch.FloatTensor) 正弦位置编码表  
        """  

        def get_position_angle_vec(i):  
            """ 获取位置角度向量 """  
            return [i / np.power(10000, 2 * (j // 2) / token_len) for j in range(token_len)]  

        sinusoid_table = np.array([get_position_angle_vec(i) for i in range(num_tokens)])  
        sinusoid_table[:, 0::2] = np.sin(sinusoid_table[:, 0::2])  
        sinusoid_table[:, 1::2] = np.cos(sinusoid_table[:, 1::2])   

        return torch.FloatTensor(sinusoid_table).unsqueeze(0)  

    PE = get_sinusoid_encoding(num_tokens+1, token_len)  
    # 计算位置嵌入 PE  
    print('位置嵌入 PE 的维度如下\n\ttoken数量:', PE.shape[1], '\n\ttoken长度:', PE.shape[2])  

    x = x + PE  
    print('加上位置嵌入后的维度如下\n\t批次大小:', x.shape[0], '\n\ttoken数量:', x.shape[1], '\n\ttoken长度:', x.shape[2])

    位置嵌入的维度是，  
       令牌数量：176 (tokens)  
       令牌长度：768  
    具有位置嵌入的维度是，  
       批大小：13  
       令牌数量：176 (tokens)  
       令牌长度：768

我们的令牌们准备好了，现在可以进入编码阶段了。

编码块是模型实际从图像令牌中学习的地方。用户可以设置编码块的数量，这是一个超参数。如图所示，下面是编码块的示意图。

编码区块

以下是一个编码块的代码。

    class Encoding(nn.Module):  

        def __init__(self,  
           dim: int,  
           num_heads: int=1,  
           hidden_chan_mul: float=4.,  
           qkv_bias: bool=False,  
           qk_scale: NoneFloat=None,  
           act_layer=nn.GELU,   
           norm_layer=nn.LayerNorm):  

            """ 编码块  

                参数:  
                    dim (int): 单个token的大小  
                    num_heads (int): MSA 中注意力头的数量  
                    hidden_chan_mul (float): 神经网络部分中隐藏通道数量的乘数  
                    qkv_bias (bool): qkv 层是否学习一个加性偏差  
                    qk_scale (NoneFloat): 用于缩放查询和键的值；  
                                        如果为 None，则查询和键按 ``head_dim ** -0.5`` 缩放  
                    act_layer (nn.modules.activation): 用于激活层的 torch 神经网络层类型  
                    norm_layer (nn.modules.normalization): 用于归一化层的 torch 神经网络层类型  
            """  

            super().__init__()  

            ## 定义层  
            self.norm1 = norm_layer(dim)  
            self.attn = Attention(dim=dim,  
                                chan=dim,  
                                num_heads=num_heads,  
                                qkv_bias=qkv_bias,  
                                qk_scale=qk_scale)  
            self.norm2 = norm_layer(dim)  
            self.neuralnet = NeuralNet(in_chan=dim,  
                                    hidden_chan=int(dim*hidden_chan_mul),  
                                    out_chan=dim,  
                                    act_layer=act_layer)  

        def forward(self, x):  
            x = x + self.attn(self.norm1(x))  
            x = x + self.neuralnet(self.norm2(x))  
            return x

_numheads、_qkvbias 和 _qkscale 参数定义了 Attention 模块的组成部分。关于视觉变压器中的注意力机制的深入探讨留待另文详述。

两个参数 _hidden_chan_mul_ 和 _act_layer_ 定义了 神经网络 模块中的组件。激活层可以是任何 torch.nn.modules.activation 中的层。稍后我们将更详细地讨论 神经网络 模块。

norm_layer，可以是任何一种torch.nn.modules.normalization层。

我们现在将依次浏览图表中的每个蓝色方块及其对应的代码。我们将使用176个token，每个长度为768。我们将使用批处理大小为13，这是因为13是质数，不会与其他参数混淆。我们将使用4个attention head，因为它能均匀地划分token长度；但在编码块中，你不会看到注意力头的维度。

    # 定义输入
    num_tokens = 176
    token_len = 768
    batch = 13
    heads = 4
    x = torch.rand(batch, num_tokens, token_len)
    print('输入的尺寸为\n\t批次大小:', x.shape[0], '\n\ttoken数:', x.shape[1], '\n\ttoken长度:', x.shape[2])

    # 定义模块
    E = Encoding(dim=token_len, num_heads=heads, hidden_chan_mul=1.5, qkv_bias=False, qk_scale=None, act_layer=nn.GELU, norm_layer=nn.LayerNorm)
    # 将模型设置为评估模式
    E.eval();

输入维度为：  
批大小：13  
令牌数：176  
令牌长度：768.

现在，我们将通过一个规范层和一个注意力模块。编码块中的注意力模块进行了参数化，以保持词的长度不变。在注意力模块之后，我们实现了第一个跳过连接。

    y = E.norm1(x)  
    print('归一化后的维度为\n\t批量大小:', y.shape[0], '\n\ttoken数量:', y.shape[1], '\n\ttoken维度:', y.shape[2])  
    y = E.attn(y)  
    print('注意力机制后的维度为\n\t批量大小:', y.shape[0], '\n\ttoken数量:', y.shape[1], '\n\ttoken维度:', y.shape[2])  
    y = y + x  
    print('跳过连接后的维度为\n\t批量大小:', y.shape[0], '\n\ttoken数量:', y.shape[1], '\n\ttoken维度:', y.shape[2])

规范化后，维度是  
       批量大小: 13   
       数量: 176   
       大小: 768  
注意力后，维度是  
       批量大小: 13   
       数量: 176   
       大小: 768  
残差链接后，维度是  
       批量大小: 13   
       数量: 176   
       大小: 768

现在，我们经过另一个标准化层，然后是神经网络模块。然后我们通过第二个分叉连接结束。

    z = E.norm2(y)  
    print('归一化后的维度是\n\t批大小:', z.shape[0], '\n\ttoken的数量:', z.shape[1], '\n\ttoken的大小:', z.shape[2])  
    z = E.neuralnet(z)  
    print('经过神经网络处理后的维度是\n\t批大小:', z.shape[0], '\n\ttoken的数量:', z.shape[1], '\n\ttoken的大小:', z.shape[2])  
    z = z + y  
    print('残差连接后的维度是\n\t批大小:', z.shape[0], '\n\ttoken的数量:', z.shape[1], '\n\ttoken的大小:', z.shape[2])

    标准化后  
       批大小：13   
       标记数量：176   
       标记维度：768  
    经过神经网络后  
       批大小：13   
       标记数量：176   
       标记维度：768  
    残差连接后  
       批大小：13   
       标记数量：176   
       标记维度：768

单个编码块的内容就是这样！由于最终维度和初始维度相同，模型可以轻松地将令牌通过多个编码块进行传递，这些编码块的数量由超参数 depth 决定。

神经网络（NN）模块是编码模块中的一个子组件。NN模块非常简单，由一个全连接层、一个激活层和另一个全连接层组成。激活层可以是任何torch.nn.modules.activation中的层，该层作为输入传递到模块。可以通过配置NN模块来改变输入的形状，或者保持形状不变。由于神经网络在机器学习中很常见，不是本文的重点，我们不会详细讲解这段代码。不过，这里展示了NN模块的代码。

    class NeuralNet(nn.Module):  
        def __init__(self,  
           in_chan: int,  
           hidden_chan: NoneFloat=None,  
           out_chan: NoneFloat=None,  
           act_layer = nn.GELU):  
            """神经网络模块

                参数:
                    in_chan (int): 输入的通道数（特征数）
                    hidden_chan (NoneFloat): 隐藏层中的通道数；如果为 None，则隐藏层的通道数与输入层相同
                    out_chan (NoneFloat): 输出的通道数；如果为 None，则输出层的通道数与输入层相同
                    act_layer(nn.modules.activation): 作为激活函数的 torch 神经网络层类
            """  

            super().__init__()  

            ## 定义各层的通道数  
            hidden_chan = hidden_chan or in_chan  
            out_chan = out_chan or in_chan  

            ## 定义各层  
            self.fc1 = nn.Linear(in_chan, hidden_chan)  
            self.act = act_layer()  
            self.fc2 = nn.Linear(hidden_chan, out_chan)  

        def forward(self, x):  
            x = self.fc1(x)  
            x = self.act(x)  
            x = self.fc2(x)  
            return x

经过编码块处理后，模型必须做的最后一件事是进行预测。如下图所示的预测阶段是ViT图的一部分。

ViT预测处理组件图示（作者绘制）

我们将查看这个过程的每一步。我们将继续使用长度为768的176个tokens。我们将采用批量大小为1来演示如何进行单个预测。批量大小大于1则表示以并行方式计算这些预测。

    # 定义一个输入  
    num_tokens = 176  
    token_len = 768  
    batch = 1  
    x = torch.rand(batch, num_tokens, token_len)  
    print('输入的维度是\n\t批次大小:', x.shape[0], '\n\ttoken个数:', x.shape[1], '\n\ttoken的长度:', x.shape[2])

输入维度如下：
批量大小：1个
令牌数：176
令牌长度：768个单位

首先，所有的tokens会经过一个规范层处理。

    norm = nn.LayerNorm(token_len)  
    x = norm(x)  
    print('经过归一化处理后，维度如下：\n\t批次大小：', x.shape[0], '\n\ttoken数量：', x.shape[1], '\n\ttoken维度：', x.shape[2])

    标准化之后，维度如下  
       批次大小：1，  
       令牌数量：1001，  
       令牌维度：768

接下来，我们将预测令牌从其余令牌中分离出来。在整个编码块过程中，预测令牌不再为零，并从输入图像中获取了信息。我们将仅用它来做最终预测。

    将预测令牌赋值为 x[:, 0]  
    打印预测令牌的长度：pred_token.shape[-1]

预测 token 的长度是 768

最后，预测标记通过_head_来进行预测。_head_通常是某种神经网络的变体，根据模型的不同而有所变化。在《一图胜千言》²这篇文章里，他们在预训练阶段使用了具有一个隐藏层的多层感知器（MLP），并在微调阶段使用了一个单层线性模型。在《Token-to-Token ViT³》这篇文章里，他们使用了一个单层线性模型作为head。这里我们继续使用一个单层线性模型。

注意，输出形状是根据学习问题的参数设定的。对于分类任务，它通常是长度为_类别数量_的向量，采用所谓的独热编码。对于回归任务，输出形状可以是任意数量的预测参数（例如整数）。这里我们使用输出形状为1，来表示单个估计的回归值。

    head = nn.Linear(token_len, 1)  
    pred = head(pred_token)  
    print('预测的长度为:', (pred.shape[0], pred.shape[1]))  
    print('预测的具体数值为:', float(pred))

预测区间：(1, 1)  
预测值：-0.5474240779876709

就这样好了！模型做出了预测！

为了创建完整的ViT模块，我们使用了上面提到的Patch Tokenization模块和VIT_BACKBONE模块。VIT_BACKBONE模块定义如下，它包含令牌处理、编码块等组件。

    class ViT_Backbone(nn.Module):  
        def __init__(self,  
                    preds: int=1,  
                    token_len: int=768,  
                    num_heads: int=1,  
                    Encoding_hidden_chan_mul: float=4.,  
                    depth: int=12,  
                    qkv_bias=False,  
                    qk_scale=None,  
                    act_layer=nn.GELU,  
                    norm_layer=nn.LayerNorm):  

            """ 视觉Transformer骨干网络
                参数:
                    preds (int): 输出的预测数量
                    token_len (int): 一个token的长度
                    num_heads (int): 多头注意力层中的注意力头数量
                    Encoding_hidden_chan_mul (float): 编码模块隐藏通道数的乘数
                    depth (int): 模型中的编码块数量
                    qkv_bias (bool): qkv是否使用偏置
                    qk_scale (float/None): 用于缩放查询和键的值；如果为None，则查询和键通过`head_dim ** -0.5`进行缩放
                    act_layer (nn.modules.activation): 用作激活的torch神经网络层类
                    norm_layer (nn.modules.normalization): 用作归一化的torch神经网络层类
            """  

            super().__init__()  

            ## 定义参数
            self.num_heads = num_heads
            self.Encoding_hidden_chan_mul = Encoding_hidden_chan_mul
            self.depth = depth

            ## 定义token处理组件
            self.cls_token = nn.Parameter(torch.zeros(1, 1, self.token_len))
            self.pos_embed = nn.Parameter(data=get_sinusoid_encoding(num_tokens=self.num_tokens+1, token_len=self.token_len), requires_grad=False)

            ## 定义编码块
            self.blocks = nn.ModuleList([Encoding(dim=self.token_len,   
                                                   num_heads=self.num_heads,  
                                                   hidden_chan_mul=self.Encoding_hidden_chan_mul,  
                                                   qkv_bias=qkv_bias,  
                                                   qk_scale=qk_scale,  
                                                   act_layer=act_layer,  
                                                   norm_layer=norm_layer)  
                 for i in range(self.depth)])

            ## 定义预测处理
            self.norm = norm_layer(self.token_len)
            self.head = nn.Linear(self.token_len, preds)

            ## 从截断的正态分布中采样类token
            timm.layers.trunc_normal_(self.cls_token, std=.02)

        def forward(self, x):  
            ## 假设x已经分词

            ## 获取批次大小
            B = x.shape[0]
            ## 拼接类别token
            x = torch.cat((self.cls_token.expand(B, -1, -1), x), dim=1)
            ## 添加位置嵌入
            x = x + self.pos_embed
            ## 依次通过每个编码块
            for blk in self.blocks:
                x = blk(x)
            ## 进行归一化
            x = self.norm(x)
            ## 对类别token进行预测
            x = self.head(x[:, 0])
            return x

从ViT Backbone 模块中，我们可以定义ViT模型。

    class ViT_Model(nn.Module): 
     def __init__(self, 
        img_size: tuple[int, int, int]=(1, 400, 100), 
        patch_size: int=50, 
        token_len: int=768, 
        preds: int=1, 
        num_heads: int=1, 
        Encoding_hidden_chan_mul: float=4., 
        depth: int=12, 
        qkv_bias=False, 
        qk_scale=None, 
        act_layer=nn.GELU, 
        norm_layer=nn.LayerNorm): 

      """ 视觉Transformer模型

       参数:
        img_size (tuple[int, int, int]): 输入图像的尺寸 (通道数, 高度, 宽度)
        patch_size (int): 方形补丁的边长
        token_len (int): 令牌的输出长度
        preds (int): 输出预测的数量
        num_heads (int): 多头注意力机制中的注意力头的数量
        Encoding_hidden_chan_mul (float): 决定编码器模块中隐藏通道（特征）数量的乘数
        depth (int): 模型中的编码器块的数量
        qkv_bias (bool): qkv层是否学习可加偏差
        qk_scale (None 或 float): 用于缩放查询和键值的值；如果为 `None`，则查询和键值将被缩放为`head_dim ** -0.5`
        act_layer (nn.modules.activation): 作为激活层使用的torch神经网络层类
        norm_layer (nn.modules.normalization): 作为正则化层使用的torch神经网络层类
      """ 
      super().__init__() 

      ## 定义参数
      self.img_size = img_size 
      C, H, W = self.img_size 
      self.patch_size = patch_size 
      self.token_len = token_len 
      self.num_heads = num_heads 
      self.Encoding_hidden_chan_mul = Encoding_hidden_chan_mul 
      self.depth = depth 

      ## 定义Patch Embedding模块
      self.patch_tokens = Patch_Tokenization(img_size, 
               patch_size, 
               token_len) 

      ## 定义ViT骨干网络
      self.backbone = ViT_Backbone(preds, 
             self.token_len, 
             self.num_heads, 
             self.Encoding_hidden_chan_mul, 
             self.depth, 
             qkv_bias, 
             qk_scale, 
             act_layer, 
             norm_layer) 
      ## 初始化权重参数
      self.apply(self._init_weights) 

     def _init_weights(self, m): 
      """ 初始化线性层和层归一化层的权重
      """ 
      ## 对于线性层
      if isinstance(m, nn.Linear): 
       ## 权重从截断正态分布中初始化
       timm.layers.trunc_normal_(m.weight, std=.02) 
       if isinstance(m, nn.Linear) and m.bias is not None: 
        ## 如果存在偏差，将偏差初始化为零
        nn.init.constant_(m.bias, 0) 
      ## 对于层归一化层
      elif isinstance(m, nn.LayerNorm): 
       ## 权重初始化为1
       nn.init.constant_(m.weight, 1.0) 
       ## 偏差初始化为零
       nn.init.constant_(m.bias, 0) 

     @torch.jit.ignore ## 告诉PyTorch不要编译为TorchScript
     def no_weight_decay(self): 
      """ 在优化器中忽略类令牌的权重衰减项
      """ 
      return {'cls_token'} 

     def forward(self, x): 
      x = self.patch_tokens(x) 
      x = self.backbone(x) 
      return x

在 ViT 模型 中，_imgsize、_patchsize 和 _tokenlen 这些参数定义了 Patch 编码 模块。

The _numheads ，_Encoding_hidden_channelmul ，_qkvbias ，_qkscale ，和 _actlayer 参数定义了 编码块 模块。_actlayer 可以是 torch.nn.modules.activation 中的任何层。depth 参数决定了模型包含多少个编码块。

_the_normlayer 参数定义了编码块模块内部和外部的归一化方式。它可以是任何可用的 torch.nn.modules.normalization 层，例如 BatchNorm 或 LayerNorm。

__init_weights_ 方法名来源于 T2T-ViT³ 代码。此方法可以删除以随机初始化所有学习到的权重和偏差。按照当前的实现，线性层的权重初始化为截断正态分布；线性层的偏差初始化为0；归一化层的权重初始化为1；归一化层的偏差初始化为0。

现在，你有了对ViT模型机制的深入了解，可以去训练模型了！下面是一些下载ViT模型代码的链接。其中一些允许更多的自定义设置。去吧，祝你训练愉快！

• 本文系列的GitHub仓库链接
• GitHub仓库 用于《一张图值16x16个词》²
  → 包含预训练模型以及微调代码，但不包含模型定义
• 在 PyTorch Image Models (timm) 中实现的ViT模型⁹
  timm.create_model('vit_base_patch16_224', pretrained=True)
• Phil Wang 的 vit-pytorch 库

本文由洛斯阿拉莫斯国家实验室批准发布，编号为LA-UR-23–33876。相关代码获批采用BSD-3开源许可证，编号O#4693。

要了解更多关于NLP环境中transformers的信息，可以参阅

• 视觉解释变压器 第 1 部分 功能概要：https://towardsdatascience.com/transformers-explained-visually-part-1-overview-of-functionality-95a6dd460452
• 视觉解释变压器 第 2 部分 工作原理详解：https://towardsdatascience.com/transformers-explained-visually-part-2-how-it-works-step-by-step-b49fa4a64f34

有关视觉Transformer的全面视频讲座，请参阅

• 视觉Transformer及其应用领域：请参见 https://youtu.be/hPb6A92LROc?si=GaGYiZoyDg0PcdSP

[1] Vaswani 等人（等人）(2017)Attention Is All You Need. (网址：https://doi.org/10.48550/arXiv.1706.03762)

[2] Dosovitskiy 等 (2020). 一张图片抵得上16x16个单词：大规模图像识别中的变压器. https://doi.org/10.48550/arXiv.2010.11929.

[3] 元等人 (2021). Tokens-to-Token ViT: 从零开始在 ImageNet 上训练视觉变压器，. https://doi.org/10.48550/arXiv.2101.11986
→ GitHub 代码: 链接: https://github.com/yitu-opensource/T2T-ViT

[4] Luis Zuno (@ansimuz)。黄昏山景背景图。CC0： https://opengameart.org/content/mountain-at-dusk-background

[5] PyTorch. 展开. https://pytorch.org/docs/stable/generated/torch.nn.Unfold.html#torch.nn.Unfold

[6] PyTorch. unsqueeze（unsqueeze）. 链接

[7] PyTorch. 非线性激活函数（加权求和，非线性）https://pytorch.org/docs/stable/nn.html#non-linear-activations-weighted-sum-nonlinearity

[8] PyTorch. 归一化层。 https://pytorch.org/docs/stable/nn.html#normalization-layers

[9] Ross Wightman. PyTorch 图像模型项目.https://github.com/huggingface/pytorch-image-models