Vision Transformer (ViT) 的核心源码实现主要围绕如何将图像转换为Transformer能够处理的序列形式，并构建完整的编码器结构。其核心流程可以解构为图像分块（Patch Embedding）、位置编码（Position Embedding）、类别标记（Class Token）的引入，以及多层Transformer Encoder的堆叠[ref_1][ref_4]。 以下将以PyTorch为例，结合关键代码片段，详细解析ViT的源码实现原理。 ### 1. 图像分块与嵌入 (Patch Embedding) 这是ViT的第一步，目的是将二维图像转换为一系列一维的向量序列。通常，将一张尺寸为 `(H, W, C)` 的图像分割成 `N` 个大小为 `P x P` 的块（patch），然后将每个块展平并线性投影到一个固定维度 `D` 的嵌入空间中[ref_1][ref_3]。 ```python import torch import torch.nn as nn class PatchEmbedding(nn.Module): def __init__(self, img_size=224, patch_size=16, in_channels=3, embed_dim=768): super().__init__() self.img_size = img_size self.patch_size = patch_size self.num_patches = (img_size // patch_size) ** 2 # 计算patch数量 N # 使用一个卷积层同时完成分块和线性投影 # 卷积核大小和步长等于patch_size，输出通道数等于嵌入维度 self.proj = nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=patch_size) def forward(self, x): # x 形状: (B, C, H, W) B, C, H, W = x.shape assert H == self.img_size and W == self.img_size, \ f"输入图像尺寸({H}*{W})与模型预设({self.img_size}*{self.img_size})不符" # 卷积操作：将图像分割为patch并投影到嵌入空间 # 输出形状: (B, embed_dim, H/patch_size, W/patch_size) x = self.proj(x) # 将空间维度展平为序列长度，并调整维度顺序 # 输出形状: (B, num_patches, embed_dim) x = x.flatten(2).transpose(1, 2) return x ``` **关键点**：使用 `nn.Conv2d` 实现分块和投影是高效的技巧，卷积核大小和步长等于 `patch_size`，使得每个卷积操作恰好处理一个patch，输出特征图的每个空间位置对应一个patch的嵌入向量[ref_1][ref_4]。 ### 2. 类别标记与位置编码 (Class Token & Position Embedding) 为了进行图像分类，ViT引入了一个可学习的 `[class]` token，其嵌入向量将作为整个图像的全局表示，送入后续的Transformer Encoder。同时，由于Transformer本身不具有位置感知能力，需要添加位置编码（Position Embedding）来保留patch之间的空间顺序信息[ref_1][ref_3]。 ```python class ViTEmbeddings(nn.Module): def __init__(self, img_size=224, patch_size=16, in_channels=3, embed_dim=768, dropout_rate=0.1): super().__init__() self.patch_embed = PatchEmbedding(img_size, patch_size, in_channels, embed_dim) self.num_patches = self.patch_embed.num_patches # 可学习的类别标记 [class] token self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) # 可学习的位置编码，长度为 (num_patches + 1)，加1是为了class token self.pos_embed = nn.Parameter(torch.zeros(1, self.num_patches + 1, embed_dim)) self.dropout = nn.Dropout(dropout_rate) # 初始化参数 nn.init.trunc_normal_(self.cls_token, std=0.02) nn.init.trunc_normal_(self.pos_embed, std=0.02) def forward(self, x): B = x.shape[0] # 批大小 # 1. 生成patch嵌入 x = self.patch_embed(x) # (B, num_patches, embed_dim) # 2. 添加class token：将可学习的cls_token复制B份，并拼接到序列开头 cls_tokens = self.cls_token.expand(B, -1, -1) # (B, 1, embed_dim) x = torch.cat((cls_tokens, x), dim=1) # (B, num_patches+1, embed_dim) # 3. 添加位置编码 x = x + self.pos_embed # 4. 应用Dropout x = self.dropout(x) return x ``` **关键点**： * **Class Token**：这是一个可学习的向量，与patch嵌入拼接后一起输入Transformer。经过多层自注意力计算后，序列第一个位置（即class token位置）的输出向量将用于最终的分类[ref_1][ref_5]。 * **位置编码**：ViT通常采用**可学习的一维位置编码**，而不是Transformer原论文中的正弦余弦编码。它为序列中的每个位置（包括class token）分配一个独特的可学习向量[ref_3][ref_6]。 ### 3. Transformer Encoder 层 ViT的核心是由多个相同的Transformer Encoder层堆叠而成。每个Encoder层包含一个**多头自注意力（Multi-Head Self-Attention, MSA）**模块和一个**前馈网络（Feed-Forward Network, FFN）**，每个模块前后都有层归一化（LayerNorm）和残差连接（Residual Connection）[ref_1][ref_2]。 | 组件 | 功能描述 | 关键实现 | | :--- | :--- | :--- | | **多头自注意力 (MSA)** | 使序列中的每个元素（patch）都能关注到所有其他元素，捕捉全局依赖关系。 | 将输入拆分为多个“头”，在每个头上独立计算注意力，最后合并。 | | **前馈网络 (FFN)** | 对每个位置的表示进行非线性变换和增强。 | 通常是一个两层MLP，中间包含GELU激活函数和Dropout。 | | **层归一化 (LayerNorm)** | 对每个样本的特征维度进行归一化，稳定训练过程。 | 应用于MSA和FFN的输入前（Pre-Norm结构）。 | | **残差连接** | 将模块的输入直接加到其输出上，缓解梯度消失，帮助训练深层网络。 | `output = module(LayerNorm(input)) + input` | ```python class MultiHeadSelfAttention(nn.Module): def __init__(self, embed_dim=768, num_heads=12, dropout_rate=0.0): super().__init__() assert embed_dim % num_heads == 0, "embed_dim 必须是 num_heads 的整数倍" self.num_heads = num_heads self.head_dim = embed_dim // num_heads self.scale = self.head_dim ** -0.5 # 缩放因子，防止点积过大 # 将Q、K、V的投影合并为一个线性层，提高效率 self.qkv = nn.Linear(embed_dim, embed_dim * 3) self.attn_dropout = nn.Dropout(dropout_rate) self.proj = nn.Linear(embed_dim, embed_dim) self.proj_dropout = nn.Dropout(dropout_rate) def forward(self, x): B, N, C = x.shape # N: 序列长度 (num_patches+1) # 生成Q、K、V，并重塑为多头形式 qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4) q, k, v = qkv[0], qkv[1], qkv[2] # 每个形状: (B, num_heads, N, head_dim) # 计算注意力分数: (Q * K^T) / sqrt(d_k) attn = (q @ k.transpose(-2, -1)) * self.scale attn = attn.softmax(dim=-1) attn = self.attn_dropout(attn) # 应用注意力权重到V上 x = (attn @ v).transpose(1, 2).reshape(B, N, C) # 输出投影 x = self.proj(x) x = self.proj_dropout(x) return x class TransformerEncoderLayer(nn.Module): def __init__(self, embed_dim=768, num_heads=12, mlp_ratio=4.0, dropout_rate=0.1): super().__init__() self.norm1 = nn.LayerNorm(embed_dim) self.attn = MultiHeadSelfAttention(embed_dim, num_heads, dropout_rate) self.dropout1 = nn.Dropout(dropout_rate) self.norm2 = nn.LayerNorm(embed_dim) hidden_dim = int(embed_dim * mlp_ratio) self.mlp = nn.Sequential( nn.Linear(embed_dim, hidden_dim), nn.GELU(), nn.Dropout(dropout_rate), nn.Linear(hidden_dim, embed_dim), nn.Dropout(dropout_rate) ) def forward(self, x): # 第一个残差块：MSA + Add x = x + self.dropout1(self.attn(self.norm1(x))) # 第二个残差块：FFN + Add x = x + self.mlp(self.norm2(x)) return x ``` **关键点**：ViT通常采用 **Pre-LayerNorm** 结构，即在MSA和FFN之前进行层归一化，这种结构被认为比原始Transformer的Post-LayerNorm更稳定，易于训练[ref_3]。 ### 4. 完整的ViT模型 将上述所有组件组合起来，并添加一个分类头（MLP Head），就构成了完整的ViT模型。 ```python class VisionTransformer(nn.Module): def __init__(self, img_size=224, patch_size=16, in_channels=3, num_classes=1000, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4.0, dropout_rate=0.1): super().__init__() self.embeddings = ViTEmbeddings(img_size, patch_size, in_channels, embed_dim, dropout_rate) # 堆叠L个Transformer Encoder层 self.encoder_layers = nn.ModuleList([ TransformerEncoderLayer(embed_dim, num_heads, mlp_ratio, dropout_rate) for _ in range(depth) ]) self.norm = nn.LayerNorm(embed_dim) # 分类头：通常只是一个线性层，有时会接一个小的MLP self.head = nn.Linear(embed_dim, num_classes) if num_classes > 0 else nn.Identity() def forward(self, x): # 1. 嵌入层 x = self.embeddings(x) # 2. 通过所有Transformer层 for layer in self.encoder_layers: x = layer(x) # 3. 对最终输出进行层归一化（仅对class token的位置） x = self.norm(x) # 4. 提取class token的输出用于分类 cls_output = x[:, 0] # 取序列的第一个位置，即class token # 5. 分类头 logits = self.head(cls_output) return logits ``` **关键点**：在模型最后，只取序列第一个位置（即`[class]` token）经过所有Encoder层和最终LayerNorm处理后的向量，作为整个图像的全局特征表示，送入分类头得到预测结果[ref_4][ref_5]。 ### 5. 模型变体与配置 ViT论文提出了几种不同规模的模型配置，主要区别在于Transformer的深度（`depth`，即层数）、嵌入维度（`embed_dim`）和注意力头数（`num_heads`）。常见的配置如下表所示[ref_5]： | 模型名称 | 层数 (L) | 隐藏层维度 (D) | MLP大小 | 注意力头数 | 参数量 | | :--- | :---: | :---: | :---: | :---: | :---: | | **ViT-Base (ViT-B/16)** | 12 | 768 | 3072 | 12 | ~86M | | **ViT-Large (ViT-L/16)** | 24 | 1024 | 4096 | 16 | ~307M | | **ViT-Huge (ViT-H/14)** | 32 | 1280 | 5120 | 16 | ~632M | 其中，“/16”表示使用的patch大小为16x16像素。更小的patch尺寸（如/8）会产生更长的序列，计算量更大，但可能带来性能提升[ref_5]。 ### 总结 ViT源码的核心思想是通过 **Patch Embedding** 将图像“词序列化”，通过引入 **Class Token** 和 **Position Embedding** 来适配分类任务和保留位置信息，并利用堆叠的 **Transformer Encoder** 层进行强大的特征学习和全局关系建模。其实现清晰地遵循了Transformer的编码器架构，是深度学习领域“用统一架构处理不同模态数据”的典范[ref_1][ref_6]。