手把手教你实现PyTorch版ViT：图像分类任务中的Transformer实战

作者：SkyXZ

CSDN：SkyXZ～-CSDN博客

博客园：SkyXZ - 博客园

• ViT论文Arxiv地址：An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale

        最近具身智能足够火热，VLM、VLA、VLN层出不穷发展迅速，而Transformer作为这些架构最重要的底座之一，得益于其强大的建模能力、良好的可扩展性与统一的结构设计，Transformer 已经成为构建多模态智能系统的事实标准。从最初的 BERT、GPT 在 NLP 中的成功，到 ViT、CLIP、RT-1 等模型在视觉和控制领域的延伸，Transformer 构筑起了统一语言、视觉乃至动作空间的桥梁。

        既然 Transformer 成为了具身智能的基础设施，那作为一名想走进机器人、走进未来的工程师，我当然也要学会它。于是我决定从最经典、最基础的 Vision Transformer（ViT）入手，一步步从原理出发，亲手用 PyTorch 复现，并整理下这一路的学习过程与思考，作为这篇博客的分享内容。如果你也对 Transformer 在视觉领域的应用感兴趣，或者正在入门具身智能相关方向，希望这篇文章能对你有所帮助！

PS：💻 项目完整代码已上传至Github：ViT_PyTorch，如果你在阅读中有任何问题、建议或错误指出，也欢迎在评论区与我讨论，我们共同进步！

一、ViT：从论文出发理解架构设计

        在正式动手复现之前，我们先从源头出发，来读一读 Vision Transformer 的原始论文：《An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale》[arXiv:2010.11929]。这是由 Google Research 于 2020 年提出的一篇具有里程碑意义的论文，它首次展示了 纯 Transformer 架构在图像分类任务上可以不依赖任何卷积模块，依然取得优秀性能。

        Transformer 原本是为了解决语言文字处理任务而提出的模型，其设计初衷是用于建模序列数据中的长距离依赖关系。在 NLP 领域中，Transformer 能够通过自注意力机制灵活地捕捉单词之间的全局关系，极大提升了语言理解与生成的能力。而谷歌的研究团队提出了非常大胆也非常优雅的一个思想：如果我们能把图像切割成小块（Patch），再把每个 Patch 当作一个“词”，是否也能将图像转化为序列，从而让 Transformer 也能处理视觉信息？而其提出的ViT 就是这样做的：它将一张图像划分为固定大小的 Patch（如 16×16），将每个 Patch 展平成向量，再通过一个线性投影层将其映射到统一的维度空间，最终形成一个 token 序列。随后，ViT 在这个 token 序列前加上一个可学习的 [CLS] token，并叠加位置编码（Positional Encoding），以保留图像中的空间位置信息。整个序列就像一段文本，送入多层标准的 Transformer 编码器结构进行处理，最后通过 CLS token 的输出，完成整张图像的分类任务。这种方法不依赖任何卷积操作，完全基于序列建模，展现了 Transformer 在图像建模上的巨大潜力。

        ViT的架构如上图，与寻常的分类网络类似，整个Vision Transformer可以分为两部分，一部分是特征提取部分，另一部分是分类部分，特征提取部分是其最核心的组成，它包括了Patch Embedding、Positional Encoding以及Transformer Encoder，分类部分 则是紧接在特征提取之后，通过一个可学习的 [CLS] token 来代表整张图像的全局语义。这个 token 会随着其他 token 一起参与 Transformer 编码过程，最终被送入一个简单的 MLP 分类头 进行类别预测。接下来我们按照如下的划分来逐个讲解ViT网络架构

• 图像分块与线性嵌入模块（Patch Embedding）

        ViT 的第一步操作，就是将输入图像转化为一系列的 视觉 token，这个过程被称为 Patch Embedding，Patch 指的就是分割后的一小块图像区域，它的核心思想非常直接：

将一张二维图像按照固定大小（如 16×16）划分成若干个小块（Patch），然后将每个 Patch 展平成一个向量，再通过一个线性层将其映射到指定的维度空间（例如 768维），从而得到一组输入 token，供 Transformer 使用。

        这个处理方式本质上就是在模拟 NLP 中“将每个单词编码为向量”的过程——只不过这里的“单词”是图像块 patch，而不是文字，我们假设假设输入图像大小为 224×224×3，Patch 大小为 16×16，则一张图像将被划分为$ (224/16)^2=14×14=196$ 个 patch，而每个 Patch 将被展平成一个 (16 × 16 × 3 = 768) 维的向量，将其展平成向量后，再通过一个 Linear 层映射到模型的 embedding 空间（手动设置，ViT-Base为 768 维，ViT-Large为1024，ViT-Huge为1280，通常使用768），最终我们就能得到一个形状为：[batch_size, 196, embed_dim]的patch token 序列，而我们该如何对图像进行分割实现Patch Embedding呢？这时候我们便可以想到我们的卷积，由于卷积使用的是滑动窗口的思想，因此我们只需要将卷积核以及步长设置成与Patch-Size相等便可，这时两个图片区域的特征提取过程就不会有重叠，当我们输入的图片是[224, 224, 3]的时候，我们可以获得一个[14, 14, 768]的特征层。

        而获得了特征信息之后我们需要将得到的特征信息组合成序列，组合的方式很简单，我们只需要对这个特征图进行展平（Flatten）并转置为标准序列格式，便可以得到最终的 Patch Token 序列，用于输入 Transformer，我们在上面对图像进行分割后得到了一个[14, 14, 768]的特征层，我们将这个特征图的高宽维度进行平铺后即可得到一个[196, 768]的特征层，至此Patch Embedding便完成啦！

• 分类标记与位置编码模块（cls_token + Position Embedding）

        在完成 Patch Embedding 得到形如 [batch_size, 196, 768] 的 Patch Token 序列后，接下来我们要做两件关键的事情：

1. 添加 [CLS] Token —— 图像的“全局摘要”入口

           Transformer 最初在处理文本任务时，会在序列的最前面添加一个特殊的 [CLS] Token，用于聚合整个句子的语义信息。同理，在 ViT 中也引入了 [CLS] Token，它并不代表某个具体的 Patch，而是作为一个全局的代表 Token，在 Transformer 中“参与”每一层的信息交互，最终用于提取整个图像的全局特征。如上图所示，编号为 0* 的那个位置即表示 [CLS] Token，其初始值是一个可学习的参数向量，维度与 Patch Token 相同（例如 768），经过 Transformer 编码后，ViT 会使用这个 [CLS] Token 的输出向量作为图像的分类结果输入到 MLP Head 中，完成最终分类。

           添加了 [CLS] Token 之后，原本的 196 个 Patch Token 序列就变成了 197 个 Token，形状变为了形如：[batch_size, 196 + 1, 768]

2. 添加位置编码（Positional Embedding）—— 帮助模型理解“图像中的位置”

           由于 Transformer 是完全基于自注意力机制构建的，它并不具备卷积网络中天然的位置信息建模能力。所以我们还需要给每个 Token 添加一个位置编码，用于告诉模型这个 Token 来自于图像的哪一块区域。ViT 采用的是一种 可学习的绝对位置编码，也就是为每一个 Token 的位置（包括 [CLS] Token）都初始化一个可学习的位置向量，并与原始 Token 相加，这样，模型就能在学习过程中自己掌握空间顺序和语义之间的关系。

           位置编码的形状与输入序列一致，也是[1, 196 + 1, 768]，且位置编码的加入方式非常简单即：tokens = tokens + pos_embed # [B, 197, 768],经过这两个步骤之后，ViT 的输入才真正准备好，可以送入 Transformer 编码器中进行多层特征交互与建模，至此cls_token + Position Embedding便完成啦！

• 标准 Transformer 编码器（Multi-head Attention + LayerNorm + MLP + 残差连接）

        当我们得到了带有 [CLS] Token 和位置编码的完整 Patch 序列（形状为 [B, 197, 768]）之后，ViT 会将其送入一系列标准的 Transformer Encoder Block 中进行深度建模。每一个 Block 的设计与原始的 NLP Transformer 中的 Encoder 保持一致，结构非常经典，由两个子模块组成：

1. LayerNorm + 多头自注意力机制（Multi-Head Self Attention）

           在这个子模块中，我们首先对输入进行 LayerNorm 归一化，再送入 Multi-Head Self-Attention 模块，这里的自注意力的作用是建立所有 Token 之间的全局关系，使每个 Token 都能获取其他区域的信息，这也是Transformer的灵魂部分，其具体实现则是让每个 Token 通过查询（Query）与所有其他 Token 的键（Key）进行匹配，计算其对其他位置的关注权重，从而提取对当前任务最有用的信息。用公式和图片表示如下：

[text{Attention}(Q, K, V) = text{Softmax}left( frac{QK^top}{sqrt{d_k}} right) V ]

        对于初学者来说，这部分内容看起来可能会比较抽象，但是我们如果将它拆解一步一步来看，其实非常直观。在多头自注意力机制中，每一个输入的 Token（图像的 Patch）都会被分别映射出三个向量，分别是：Query（查询向量）、Key（键向量）、Value（值向量），用一个简单明了的比喻来理解：假设你在参加一次会议（注意力机制），你是 Query，而会议室里每一个与会者（包括你自己）都是一个 Key，同时他们手里都拿着一份资料（Value），你会根据自己和其他人 Key 的“相似程度”决定你要多大程度参考他们的资料（Value）——这就是注意力权重的计算。设当前输入序列为矩阵 (X in mathbb{R}^{n times d})，其中 (n) 是序列长度（例如 ViT 中是 197 个 Token），(d) 是每个 Token 的维度（例如 768）。我们用三组可学习的参数矩阵将其变换为：

[[ Q = XW^Q,quad K = XW^K,quad V = XW^V ] ]

        然后计算注意力得分（Score）：

[[ text{Score} = frac{QK^top}{sqrt{d_k}} ] ]

        接着使用 Softmax 对得分进行归一化，得到注意力权重 (alpha)：

[[ alpha = text{Softmax}left( frac{QK^top}{sqrt{d_k}} right) ] ]

        最后加权组合所有 Value 向量，得到新的输出表示：

[[ text{Attention}(Q, K, V) = alpha V ] ]

        这套计算流程用人话说就是，我们这套注意力系统假设有三个输入分别是input-1、input-2、input-3以及三个对应的输出output-1、output-2、output-3，每个输入都有他们自己的QKV向量

        如果我们要求output-1，那我们首先先将input-1的Q查询向量与三个输入的K键向量分别相乘得到对应的分数，这个分数代表的便是 input-1 对其它三个输入的“注意力程度”；接下来我们将这三个分数分别求一次softmax使它们变成 0 到 1 之间的概率值，并且加起来为 1，这个过程可以理解为：分配关注度，告诉我们该“关注谁、关注多少”；最后，用刚才得到的这三个注意力权重，去分别加权对应的 值向量 V，再把它们加在一起，得到的就是最终的 output-1

        也就是说：输出 = 所有关注对象的“值” × “关注它的程度”的加权和。每个 Query 会根据与所有 Key 的相似程度，对对应的 Value 进行加权求和；这样的话所有 Token 之间都能进行信息交换，从而捕捉 全局上下文依赖；因此最后的输出的虽然仍然是一个与原始 Token 数量相同的新序列，但每个 Token 的表示已经融合了全局信息。

2. LayerNorm + MLP 前馈神经网络

           在每个 Transformer Block 中，除了注意力机制之外，还有一个非常重要的部分，那就是 前馈神经网络（Feed Forward Network, FFN），也常被称为 MLP 子模块。这个子模块的结构其实非常简单，就是两个全连接层（Linear），中间再加一个非线性激活函数（如 GELU）：

   [FFN(x)=Linear 2 ​ (GELU(Linear 1 ​ (x))) ]

           这里的 Linear 层也就是我们熟悉的全连接层，维度的变化一般是这样的：首先第一个 Linear 层会把输入的维度从 d_model（比如 768）提升到一个更高的维度（比如 3072），接着通过 GELU 激活函数引入非线性，最后再用一个 Linear 层将维度降回原来的 d_model这个 FFN 的结构可以理解为对每个 Token 独立地进行更深层次的特征变换。不同于多头自注意力机制那种跨 Token 的信息交互，前馈神经网络的处理是逐 Token 的点对点非线性变换，主要用于增强模型的表达能力。而残差连接的引入可以在每个 Transformer Block 内形成一种短路路径（Shortcut Path），它能有效缓解深层网络中的梯度消失问题：与其直接学习一个映射函数 (F(x))，不如让网络学习 (F(x) = H(x) - x)，即让模型关注“输入与输出的差值”，这样反而更容易优化。因此完整的计算流程如下：

   [y=x+FFN(LayerNorm(x)) ]

        至此，ViT 的核心结构也就完整拼装完成了。从 Patch Embedding 到 [CLS] Token 与位置编码，再到深度的多层 Transformer 编码器，ViT 完整地将语言模型的结构移植到了视觉领域，并取得了突破性的表现。Transformer Block 是 ViT 的“建模大脑”，也是其通用性与强大性能的根基。

• 分类头（Classification Head）

        经过多个 Transformer Block 的深度特征提取之后，我们得到了一个新的序列表示，其形状为 [B, 197, 768]（假设我们使用的是 ViT-Base 模型），其中第一个位置的 Token 仍然是我们在最开始加入的 [CLS] Token。这个 [CLS] Token 可以看作是整个图像的全局语义表示，因为在多轮注意力交互中，它已经“融合”了所有 Patch 的信息。因此，我们只需要从序列中取出这一位置的向量（即第一个 Token），然后送入一个全连接层（Linear）就可以完成分类任务了。

二、实战复现PyTorch版ViT网络架构

（一）模块1：PatchEmbedding类

        PatchEmbedding 是 ViT 中最关键的一步，我们在这里使用卷积操作将输入图像划分为若干不重叠的小块（Patch），每个 Patch 被编码为一个向量，我们采用等步长卷积的方式实现划分，并在展平后将其送入 Transformer 模块进行后续处理。

class VisionPatchEmbedding(nn.Module):     def __init__(self, image_size, patch_size, in_channels, embed_dim, flatter=True):         super().__init__()         self.proj = nn.Conv2d(in_channels, embed_dim, patch_size, patch_size)         self.norm = nn.LayerNorm(embed_dim)         self.flatter = flatter      def forward(self, x):         x = self.proj(x)         if self.flatter:             x = x.flatten(2).transpose(1, 2)  # [B, C, H, W] -> [B, N, C]         x = self.norm(x)         return x 

（二）模块2：PositionEmbedding

        由于ViT 不像CNN，其没有卷积感受野，因此需要加入位置编码（pos_embed）来保留位置信息每一个Patch的位置信息

self.cls_token = nn.Parameter(torch.zeros(1, 1, num_features)) self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, num_features)) 

        由ViTViT在预训练时通常使用固定分辨率（如 224x224），图像被分割为 14x14 个 Patch（patch_size=16），位置编码 pos_embed 的形状为 [1, 197, 768]（197 = 1(cls_token) + 14x14），但是同样由于ViT没有感受野一说，因此实际应用当输入分辨率不同（如 256x256）时，Patch 数量变为 16x16 = 256（+1 cls_token = 257），原来的位置编码（197）便无法直接使用，因此我们需要通过双三次插值（bicubic）将 14x14 的位置编码调整到新分辨率对应的网格尺寸，具体实现如下：

img_token_pos_embed = F.interpolate(     img_token_pos_embed, size=self.features_shape, mode='bicubic', align_corners=False ) pos_embed = torch.cat((cls_token_pos_embed, img_token_pos_embed), dim=1) x = self.pos_drop(x + pos_embed) 

（三）模块3：Multi-head Attention与MLP

        接着我们实现Transformer中最关键的多头注意力机制，我们定义一个类SelfAttention 来实现了多头自注意力机制，在这个模块里我们先用一个线性层同时生成查询（Q）、键（K）、值（V），并按头数拆分维度，然后计算 Q 和 K 的点积并缩放，通过 softmax得到注意力权重，利用权重加权值（V），最后将多头结果拼接后通过线性变换和 dropout，使输出具有与输入维度相同的特征，完成信息的动态融合与表达增强。

class SelfAttention(nn.Module):     def __init__(self, dim, num_heads, qkv_bias=False, attn_drop_rate=0.0, proj_drop_rate=0.0):         super().__init__()         self.num_heads = num_heads         self.head_dim = dim // num_heads         self.scale = self.head_dim ** -0.5          self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)         self.attn_drop = nn.Dropout(attn_drop_rate)         self.proj = nn.Linear(dim, dim)         self.proj_drop = nn.Dropout(proj_drop_rate)      def forward(self, x):         B, N, C = x.shape         qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2,0,3,1,4)         q, k, v = qkv[0], qkv[1], qkv[2]          attn = torch.matmul(q, k.transpose(-2, -1)) * self.scale         attn = attn .softmax(dim=-1)         attn = self.attn_drop(attn)          x = torch.matmul(attn, v).transpose(1,2).reshape(B,N,C)         x = self.proj(x)         x = self.proj_drop(x)          return x 

        在这个 MLP 模块中，我设计了一个两层的全连接网络，首先通过 fc1 将输入特征映射到隐藏维度，然后经过激活函数非线性变换，接着用 dropout 做正则化防止过拟合，接着再通过 fc2 映射到输出维度，最后再用一次 dropout，这个过程用来帮助模型捕捉更丰富的非线性特征，提升表达能力。

class MLP(nn.Module):     def __init__(self, in_features, hidden_features, out_features, act_layer, drop_rate):         super().__init__()         out_features = out_features or in_features         hidden_features = hidden_features or in_features         drop_probs = (drop_rate, drop_rate)          self.fc1 = nn.Linear(in_features, hidden_features)         self.act = act_layer()         self.drop1 = nn.Dropout(drop_probs[0])         self.fc2 = nn.Linear(hidden_features, out_features)         self.drop2 = nn.Dropout(drop_probs[1])      def forward(self, x):         x = self.fc1(x)         x = self.act(x)         x = self.drop1(x)         x = self.fc2(x)         x = self.drop2(x)         return x 

（四）模块4：Encoder层堆叠

        由于PyTorch中没有现成的DropPath函数可以使用，因此我们需要自己实现这一用法，在这里我们使用DropPath来随机丢弃完整路径实现深度网络的正则化，并在训练时以概率drop_path跳过当前模块并缩放剩余路径以保持期望值；在Block类中我则封装了完整的Transformer层结构，包含LayerNorm归一化、多头注意力、MLP前馈网络和残差连接，其中注意力部分使用我自定义的SelfAttention模块，MLP采用先扩展后压缩的结构设计，两者都集成了DropPath机制

        于是一个完整的Transformer Block计算流程如下：

graph TD A[输入 x] A --> B[LayerNorm] B --> C[Multi-Head Self-Attention] C --> D["Residual Add: x + Attention"] D --> E[LayerNorm] E --> F[FeedForward MLP] F --> G["Residual Add: D + MLP"] G --> H[输出 y]

class DropPath(nn.Module):     def __init__(self, drop_prob=None):         super(DropPath, self).__init__()         self.drop_prob = drop_prob      def drop_path(self, x, drop_prob, training):         if drop_prob == 0. or not training:             return x         keep_prob       = 1 - drop_prob         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         return output      def forward(self, x):         return self.drop_path(x, self.drop_prob, self.training)      class Block(nn.Module):     def __init__(self, dim, num_heads, mlp_radio, qkv_bias, drop, attn_drop, drop_path, act_layer, norm_layer):         super().__init__()         self.norm_1 = norm_layer(dim)         self.attn = SelfAttention(dim, num_heads=num_heads, qkv_bias=qkv_bias, attn_drop_rate=attn_drop, proj_drop_rate=drop)         self.norm_2 = norm_layer(dim)         self.mlp = MLP(in_features=dim, hidden_features=int(dim * mlp_radio), out_features=None, act_layer=act_layer, drop_rate=drop_path)         self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() # 丢弃路径      def forward(self, x):         x = x + self.drop_path(self.attn(self.norm_1(x)))         x = x + self.drop_path(self.mlp(self.norm_2(x)))         return x 

（五）模块5：ViT整体模型类

        最后我们来实现我们完整的ViT—VisonTransformer，在这个 VisonTransformer 类中，我将前面介绍的各个模块整合在一起，实现了完整的ViT网络。首先，我用卷积将输入图像切分成固定大小的Patch，并映射到特征空间；接着通过引入可学习的分类token和位置编码，为模型提供位置信息弥补卷积“感受野”缺失的问题。之后，我堆叠多个Transformer编码器Block，每个Block包含多头自注意力机制和MLP模块，通过残差连接和归一化保证信息的有效传递与特征抽象。最后，我取分类token的输出，通过线性层映射到目标类别，实现图像分类任务。

class VisonTransformer(nn.Module):     def __init__(self, input_shape, patch_size, in_channels, num_classes, num_features, depth,                  num_heads, mlp_ratio, qkv_bias, drop_rate, attn_drop_rate, drop_path_rate,                  norm_layer, act_layer):         super().__init__()         self.input_shape = input_shape # 输入的维度         self.patch_size = patch_size # Patch 的大小         self.in_channels = in_channels # 输入的维度         self.num_classes = num_classes # 输出类别数         self.num_features = num_features # 特征维度         self.depth = depth # Transformer编码器层数         self.num_heads = num_heads # Transformer注意力头数         self.mlp_ratio = mlp_ratio # MLP 比例 MLP:多层感知机,紧随 Self-Attention 之后，用于非线性变换：增强模型的表达能力；特征映射：将 Self-Attention 提取的特征进一步转换。          self.qkv_bias = qkv_bias # 是否使用偏置         self.drop_rate = drop_rate # 丢弃率         self.attn_drop_rate = attn_drop_rate # 注意力丢弃率         self.drop_path_rate = drop_path_rate # 丢弃路径率         self.norm_layer = norm_layer # 归一化层         self.act_layer = act_layer # 激活函数层          self.features_shape = [input_shape[1] // patch_size, input_shape[2] // patch_size]  # [14, 14]         self.num_patches = self.features_shape[0] * self.features_shape[1]         self.patch_embed = VisionPatchEmbedding(input_shape, patch_size, in_channels, num_features) # 将输入图片分割成patch，并进行线性映射          # ViT 不是 CNN，没有"感受野"，所以引入了位置编码，来为每个 patch 加上位置信息；         self.pretrained_features_shape = [224 // patch_size, 224 // patch_size] # 预训练的特征图尺寸          self.cls_token = nn.Parameter(torch.zeros(1, 1, num_features)) # 分类 token 196, 768 -> 197, 768         self.pos_embed = nn.Parameter(torch.zeros(1, self.num_patches + 1, num_features)) # 位置编码 197, 768 -> 197, 768          self.pos_drop = nn.Dropout(drop_rate) # 丢弃率         self.norm = norm_layer(self.num_features) # 归一化          self.dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)] # 丢弃路径率         self.blocks = nn.Sequential(             *[                 Block(                     dim = num_features,                     num_heads = num_heads,                     mlp_radio = mlp_ratio,                     qkv_bias = qkv_bias,                     drop = drop_rate,                     attn_drop = attn_drop_rate,                     drop_path = self.dpr[i],                     norm_layer = norm_layer,                     act_layer = act_layer                  )for i in range(depth)             ]         )         self.head = nn.Linear(num_features, num_classes) if num_classes > 0 else nn.Identity()      def forward_features(self,x):         x = self.patch_embed(x)         cls_token = self.cls_token.expand(x.shape[0], -1, -1) # 将分类 token 扩展到与输入特征图相同的形状         x = torch.cat((cls_token, x), dim=1) # 将分类 token 与输入特征图拼接          cls_token_pos_embed = self.pos_embed[:, 0:1, :] # 分类 token 的位置编码         img_token_pos_embed = self.pos_embed[:, 1:, :]  # [1, num_patches, num_features]         # 变成[1, H, W, C]         img_token_pos_embed = img_token_pos_embed.view(1, self.features_shape[0], self.features_shape[1], -1).permute(0, 3, 1, 2)  # [1, C, H, W]         # 插值         img_token_pos_embed = F.interpolate(             img_token_pos_embed,             size=self.features_shape,  # [H, W]             mode='bicubic',             align_corners=False         )         # 变回[1, num_patches, C]         img_token_pos_embed = img_token_pos_embed.permute(0, 2, 3, 1).reshape(1, -1, img_token_pos_embed.shape[1])          pos_embed = torch.cat((cls_token_pos_embed, img_token_pos_embed), dim=1) # 将分类 token 的位置编码与图像 token 的位置编码拼接                  x = self.pos_drop(x + pos_embed) # 将位置编码与输入特征图相加          x = self.blocks(x)         x = self.norm(x)          return x[:, 0] # 返回分类 token 的特征          def forward(self, x):         x = self.forward_features(x)         x = self.head(x)         return x 

至此我们完整的ViT便搭建完成了！

（六）实现数据加载代码（数据加载、loss、优化器）

        数据集加载部分比较简答，不过多赘述，我的数据集结构及具体代码如下：

import os  from torch.utils.data import Dataset, DataLoader  from PIL import Image  import torchvision.transforms as transforms   class ViTDataset(Dataset):     def __init__(self, root, split, transform=None, target_transform=None, img_size=224):         super().__init__()         self.split = split          self.img_size = img_size  # 图像大小         self.transform = transform if transform is not None else transforms.ToTensor()         self.target_transform = target_transform  # 标签变换         # 构建数据集根目录         self.data_dir = os.path.join(root, split)  # 训练集或测试集目录         # 获取所有类别         self.classes = sorted(os.listdir(self.data_dir))         self.class_to_idx = {cls_name: i for i, cls_name in enumerate(self.classes)}         # 收集所有图像文件路径和对应的标签         self.images = []         self.labels = []         for class_name in self.classes:             class_dir = os.path.join(self.data_dir, class_name)             if not os.path.isdir(class_dir):                 continue             for img_name in os.listdir(class_dir):                 if img_name.endswith(('.jpg', '.jpeg', '.png')):                     img_path = os.path.join(class_dir, img_name)                     self.images.append(img_path)                     self.labels.append(self.class_to_idx[class_name])         print(f"加载了 {len(self.images)} 张图像用于{split}集，共{len(self.classes)}个类别")      def __len__(self):         return len(self.images)          def __getitem__(self, index):         # 获取图像路径和标签         img_path = self.images[index]         label = self.labels[index]         # 加载图像         image = Image.open(img_path).convert('RGB')         # 调整图像大小         image = image.resize((self.img_size, self.img_size), Image.Resampling.BILINEAR)         # 应用变换         image = self.transform(image)         if self.target_transform is not None:             label = self.target_transform(label)                      return image, label       def ViTDataLoad(root, batch_size, num_workers, img_size):     # 创建训练数据集     train_dataset = ViTDataset(         root=root,         split='train',  # 使用训练集划分         img_size=img_size     )          # 创建验证数据集     val_dataset = ViTDataset(         root=root,         split='val',  # 使用验证集划分         img_size=img_size     )          # 创建训练数据加载器     train_loader = DataLoader(         train_dataset,         batch_size=batch_size,         shuffle=True,  # 随机打乱数据         num_workers=num_workers,  # 多线程加载         pin_memory=True,  # 数据预加载到固定内存，加速GPU传输         drop_last=True  # 丢弃最后不足一个批次的数据     )          # 创建验证数据加载器     val_loader = DataLoader(         val_dataset,         batch_size=batch_size,         shuffle=False,  # 不打乱数据         num_workers=num_workers,         pin_memory=True     )          return train_loader, val_loader 

（七）实现训练代码

        接下来我们来完成我们的训练代码，我通过上一节定义的ViTDataLoad来加载训练和验证数据集，训练中我采用常用的交叉熵损失函数（CrossEntropyLoss）来衡量分类效果，优化器使用了更适合Transformer的AdamW，具体实现如下，不过多赘述：

from model.transformer_net import VisonTransformer from dataset_load import ViTDataLoad import torch  import torch.nn as nn import torch.optim as optim from tqdm import tqdm import matplotlib.pyplot as plt  def train(     root="/home/xq/Working/dockertrain_test/input/timmdataset/african-wildlife", # 数据集根目录     img_size=224,     patch_size=16,     in_channels=3,     num_features=768,     depth=12,     num_heads=12,     mlp_ratio=4.0,     qkv_bias=True,     drop_rate=0.1,     attn_drop_rate=0.1,     drop_path_rate=0.1,     epochs=50,     batch_size=4,     num_workers=4,     lr=1e-4,     device=None ):     device = device or ("cuda" if torch.cuda.is_available() else "cpu")     # 数据加载     train_loader, val_loader = ViTDataLoad(root, batch_size, num_workers, img_size)     num_classes = len(train_loader.dataset.classes)     input_shape = (in_channels, img_size, img_size)      # 模型     model = VisonTransformer(         input_shape=input_shape,         patch_size=patch_size,         in_channels=in_channels,         num_classes=num_classes,         num_features=num_features,         depth=depth,         num_heads=num_heads,         mlp_ratio=mlp_ratio,         qkv_bias=qkv_bias,         drop_rate=drop_rate,         attn_drop_rate=attn_drop_rate,         drop_path_rate=drop_path_rate,         norm_layer=nn.LayerNorm,         act_layer=nn.GELU     ).to(device)      criterion = nn.CrossEntropyLoss()     optimizer = optim.AdamW(model.parameters(), lr=lr)     best_acc = 0     train_loss_list, val_loss_list = [], []     train_acc_list, val_acc_list = [], []      for epoch in range(epochs):         model.train()         total_loss, correct, total = 0, 0, 0         pbar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}")         for images, labels in pbar:             images, labels = images.to(device), labels.to(device)             optimizer.zero_grad()             outputs = model(images)             loss = criterion(outputs, labels)             loss.backward()             optimizer.step()             total_loss += loss.item() * images.size(0)             _, preds = outputs.max(1)             correct += preds.eq(labels).sum().item()             total += labels.size(0)         train_loss = total_loss / total         train_acc = correct / total         train_loss_list.append(train_loss)         train_acc_list.append(train_acc)          # 验证         model.eval()         val_loss, val_correct, val_total = 0, 0, 0         with torch.no_grad():             for images, labels in val_loader:                 images, labels = images.to(device), labels.to(device)                 outputs = model(images)                 loss = criterion(outputs, labels)                 val_loss += loss.item() * images.size(0)                 _, preds = outputs.max(1)                 val_correct += preds.eq(labels).sum().item()                 val_total += labels.size(0)         val_loss = val_loss / val_total         val_acc = val_correct / val_total         val_loss_list.append(val_loss)         val_acc_list.append(val_acc)          print(f"Epoch {epoch+1}: Train Loss={train_loss:.4f}, Train Acc={train_acc:.4f}, Val Loss={val_loss:.4f}, Val Acc={val_acc:.4f}")         # 保存最优模型         if val_acc > best_acc:             best_acc = val_acc             torch.save(model.state_dict(), "best_vit.pth")      # 可视化loss和acc     plt.figure()     plt.plot(train_loss_list, label="Train Loss")     plt.plot(val_loss_list, label="Val Loss")     plt.legend()     plt.title("Loss Curve")     plt.savefig("loss_curve.png")     plt.figure()     plt.plot(train_acc_list, label="Train Acc")     plt.plot(val_acc_list, label="Val Acc")     plt.legend()     plt.title("Accuracy Curve")     plt.savefig("acc_curve.png")     print("训练完成，最优验证准确率：", best_acc)  if __name__ == "__main__":     train() 

（八）实现验证代码

        验证代码也比较简单，有PyTorch及深度学习基础的同学可以很快实现，故这里也不再赘述：

import torch from model.transformer_net import VisonTransformer import torchvision.transforms as transforms from PIL import Image import sys import os  img_size = 224 patch_size = 16 in_channels = 3 num_features = 768 depth = 12 num_heads = 12 mlp_ratio = 4.0 qkv_bias = True drop_rate = 0.1 attn_drop_rate = 0.1 drop_path_rate = 0.1  classes = ['cat', 'dog']   num_classes = len(classes) input_shape = (in_channels, img_size, img_size)  def load_model(device):     model = VisonTransformer(         input_shape=input_shape,         patch_size=patch_size,         in_channels=in_channels,         num_classes=num_classes,         num_features=num_features,         depth=depth,         num_heads=num_heads,         mlp_ratio=mlp_ratio,         qkv_bias=qkv_bias,         drop_rate=drop_rate,         attn_drop_rate=attn_drop_rate,         drop_path_rate=drop_path_rate,         norm_layer=torch.nn.LayerNorm,         act_layer=torch.nn.GELU     ).to(device)     model.load_state_dict(torch.load("best_vit.pth", map_location=device))     model.eval()     return model  def predict(img_path, model, device):     transform = transforms.Compose([         transforms.Resize((img_size, img_size)),         transforms.ToTensor(),     ])     img = Image.open(img_path).convert('RGB')     img = transform(img).unsqueeze(0).to(device)     with torch.no_grad():         output = model(img)         pred = output.argmax(dim=1).item()     return classes[pred]  if __name__ == "__main__":     img_path = sys.argv[1]     if not os.path.exists(img_path):         print(f"图片不存在: {img_path}")          sys.exit(1)     device = "cuda" if torch.cuda.is_available() else "cpu"     model = load_model(device)     pred_class = predict(img_path, model, device)     print(f"图片 {img_path} 的预测类别为: {pred_class}")  

三、ViT：在自定义数据集及CIFAR-10进行训练与测试

（一）自定义数据集

        我们首先在我们自定义数据集上进行测试，请保证数据集格式为：

        接着在训练代码中修改数据集路径并运行下述命令即可开始训练：

python3 train.py 

        训练完成后文件内会有一个best_vit.pth以及两个训练的Acc及Loss图用于分析

接下来我们运行python3 predict.py [img_path]即可执行推理啦！

（二）CIFAR-10数据集

        在完成上述自定义数据集之后我们便可以继续尝试CIFAR-10啦！我们首先来下载我们的CIFAR-10数据集，CIFAR-10数据集已经集成进了Torch，因此我们可以使用PyTorch接口直接下载，具体下载方式如下，不过多赘述：

# CIFAR-10全量 import torchvision.datasets as datasets train_dataset = datasets.CIFAR10(root='./data', train=True, download=True) test_dataset = datasets.CIFAR10(root='./data', train=False, download=True) 

        由于我们下载下来的CIFAR-10数据集格式如下（data_batch_1 ~ data_batch_5：训练数据（每个10,000张图像），test_batch：测试数据（10,000张图像），batches.meta：元数据文件（包含类别名称等信息）），因此我们需要对我们的dataset_load.py代码进行修改以适配我们的CIFAR-10数据集

import os import pickle import numpy as np import torch from torch.utils.data import Dataset, DataLoader from PIL import Image import torchvision.transforms as transforms  class CIFAR10Dataset(Dataset):     def __init__(self, root, train=True, transform=None, target_transform=None):         super().__init__()         self.root = root         self.train = train         self.transform = transform         self.target_transform = target_transform                  # CIFAR-10类别名称         self.classes = [             'airplane', 'automobile', 'bird', 'cat', 'deer',             'dog', 'frog', 'horse', 'ship', 'truck'         ]         self.class_to_idx = {cls_name: i for i, cls_name in enumerate(self.classes)}                  self.data = []         self.targets = []                  if self.train:             for i in range(1, 6):                 batch_file = os.path.join(root, f'data_batch_{i}')                 with open(batch_file, 'rb') as f:                     batch_data = pickle.load(f, encoding='bytes')                     self.data.append(batch_data[b'data'])                     self.targets.extend(batch_data[b'labels'])             self.data = np.vstack(self.data)         else:             test_file = os.path.join(root, 'test_batch')             with open(test_file, 'rb') as f:                 test_data = pickle.load(f, encoding='bytes')                 self.data = test_data[b'data']                 self.targets = test_data[b'labels']                  # 将数据reshape为图像格式 (N, 32, 32, 3)         self.data = self.data.reshape(-1, 3, 32, 32).transpose(0, 2, 3, 1)                  print(f"加载了 {len(self.data)} 张CIFAR-10图像用于{'训练' if train else '测试'}，共{len(self.classes)}个类别")      def __len__(self):         return len(self.data)          def __getitem__(self, index):         img = self.data[index]         target = self.targets[index]         img = Image.fromarray(img)         if self.transform is not None:             img = self.transform(img)                  if self.target_transform is not None:             target = self.target_transform(target)                      return img, target  def CIFAR10DataLoad(root, batch_size, num_workers=4, img_size=224):     train_transform = transforms.Compose([         transforms.Resize((img_size, img_size)),           transforms.RandomHorizontalFlip(p=0.5),          transforms.RandomRotation(10),          transforms.ToTensor(),         transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])      ])          test_transform = transforms.Compose([         transforms.Resize((img_size, img_size)),         transforms.ToTensor(),         transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])     ])          # 创建数据集     train_dataset = CIFAR10Dataset(         root=root,         train=True,         transform=train_transform     )          test_dataset = CIFAR10Dataset(         root=root,         train=False,         transform=test_transform     )          # 创建数据加载器     train_loader = DataLoader(         train_dataset,         batch_size=batch_size,         shuffle=True,         num_workers=num_workers,         pin_memory=True,         drop_last=True     )          test_loader = DataLoader(         test_dataset,         batch_size=batch_size,         shuffle=False,         num_workers=num_workers,         pin_memory=True     )          return train_loader, test_loader   if __name__ == "__main__":     # 测试数据加载器     root = "/home/xq/Temp/cifar-10-batches-py"     train_loader, test_loader = CIFAR10DataLoad(root, batch_size=32)     # 测试一个batch     for images, labels in train_loader:         print(f"图像batch形状: {images.shape}")         print(f"标签batch形状: {labels.shape}")         print(f"标签范围: {labels.min()} - {labels.max()}")         break  

        接着修改我们训练代码，主要修改数据集路径、数据加载器的调用即可：

python3 train.py 

        由于CIFAR-10数据集比较大，训练速度较慢，我们耐心等待即可，训练完成后运行如下命令即可执行推理：

# 单个预测 python3 predict_cifar10.py <图片路径> # Top-K预测 python3 predict_cifar10.py <图片路径> --top-k 3