Vision Transformer (ViT) 图像分类实战：PyTorch 代码逐行解析与 12 层 Encoder 调参

Vision Transformer (ViT) 图像分类实战：PyTorch 代码逐行解析与 12 层 Encoder 调参

在计算机视觉领域，Transformer架构正逐渐取代传统的卷积神经网络（CNN），成为图像处理的新范式。本文将深入探讨Vision Transformer（ViT）的核心机制，并通过PyTorch实现一个完整的12层Encoder模型，同时提供关键参数调整的实验策略。

1. ViT核心架构解析

1.1 图像分块嵌入（Patch Embedding）

ViT的核心创新在于将图像视为序列数据。标准实现中，224×224的输入图像被分割为16×16的块（patch），每个patch大小为14×14。通过卷积操作实现这一过程：

class PatchEmbed(nn.Module): def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768): super().__init__() self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size) def forward(self, x): x = self.proj(x) # [B, C, H, W] -> [B, E, H/P, W/P] x = x.flatten(2).transpose(1, 2) # [B, E, N] -> [B, N, E] return x

关键参数说明：

• embed_dim：每个patch的嵌入维度（默认768）
• patch_size：决定分割粒度，影响计算复杂度和模型性能

1.2 位置编码与类别标记

与传统Transformer不同，ViT需要处理二维空间信息：

# 可学习的位置编码 self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim)) # 类别标记（CLS Token） self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))

位置编码的几种实现方式对比：

1.3 Transformer Encoder结构

标准ViT的Encoder由多层相同结构堆叠而成，每层包含：

class TransformerBlock(nn.Module): def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, drop=0.): super().__init__() self.norm1 = nn.LayerNorm(dim) self.attn = MultiHeadAttention(dim, num_heads, qkv_bias) self.norm2 = nn.LayerNorm(dim) self.mlp = MLP(dim, int(dim*mlp_ratio)) def forward(self, x): x = x + self.attn(self.norm1(x)) x = x + self.mlp(self.norm2(x)) return x

2. 关键模块实现细节

2.1 多头注意力机制

class MultiHeadAttention(nn.Module): def __init__(self, dim, num_heads=8, qkv_bias=False): super().__init__() self.num_heads = num_heads head_dim = dim // num_heads self.scale = head_dim**-0.5 self.qkv = nn.Linear(dim, dim*3, bias=qkv_bias) self.proj = nn.Linear(dim, dim) def forward(self, x): B, N, C = x.shape qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C//self.num_heads) q, k, v = qkv.permute(2, 0, 3, 1, 4) attn = (q @ k.transpose(-2, -1)) * self.scale attn = attn.softmax(dim=-1) x = (attn @ v).transpose(1, 2).reshape(B, N, C) x = self.proj(x) return x

提示：注意力头的数量需要能被嵌入维度整除，典型配置为12头×64维（总计768维）

2.2 前馈网络（MLP）

class MLP(nn.Module): def __init__(self, in_features, hidden_features=None): super().__init__() hidden_features = hidden_features or in_features self.fc1 = nn.Linear(in_features, hidden_features) self.act = nn.GELU() self.fc2 = nn.Linear(hidden_features, in_features) def forward(self, x): return self.fc2(self.act(self.fc1(x)))

扩展比例（mlp_ratio）是重要超参数：

• 基准模型：ratio=4（3072维中间层）
• 大型模型：ratio=8（6144维中间层）

3. 完整ViT模型实现

3.1 模型架构代码

class VisionTransformer(nn.Module): def __init__(self, img_size=224, patch_size=16, in_chans=3, num_classes=1000, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4., qkv_bias=True): super().__init__() self.patch_embed = PatchEmbed(img_size, patch_size, in_chans, embed_dim) num_patches = (img_size // patch_size) ** 2 self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim)) self.blocks = nn.Sequential(*[ TransformerBlock(embed_dim, num_heads, mlp_ratio, qkv_bias) for _ in range(depth)]) self.norm = nn.LayerNorm(embed_dim) self.head = nn.Linear(embed_dim, num_classes) def forward(self, x): x = self.patch_embed(x) cls_token = self.cls_token.expand(x.shape[0], -1, -1) x = torch.cat((cls_token, x), dim=1) x = x + self.pos_embed x = self.blocks(x) x = self.norm(x[:, 0]) # 仅取CLS token return self.head(x)

3.2 参数初始化策略

def init_weights(m): if isinstance(m, nn.Linear): nn.init.trunc_normal_(m.weight, std=0.02) if m.bias is not None: nn.init.zeros_(m.bias) elif isinstance(m, nn.LayerNorm): nn.init.ones_(m.weight) nn.init.zeros_(m.bias) model.apply(init_weights)

4. 关键参数调优实验

4.1 Encoder层数影响

通过改变depth参数测试不同层数的表现：

注意：超过12层后性能提升边际效应明显，需权衡计算成本

4.2 注意力头数量调整

# 不同头数配置实验 configs = [ {'num_heads': 6, 'head_dim': 128}, # 总维度768 {'num_heads': 12, 'head_dim': 64}, # 基准配置 {'num_heads': 16, 'head_dim': 48} # 非整除配置（不推荐） ]

实验发现：

• 头数增加能提升模型容量，但会降低每个头的表征能力
• 头维度应保持在64左右为佳

4.3 学习率调度策略

推荐使用余弦退火配合warmup：

optimizer = AdamW(model.parameters(), lr=1e-3) scheduler = CosineAnnealingLR( optimizer, T_max=100, eta_min=1e-5 ) warmup = GradualWarmupScheduler( optimizer, multiplier=1, total_epoch=10 )

典型训练曲线：

1. 前10个epoch线性warmup到3e-4
2. 后续90个epoch余弦退火到1e-5
3. 批量大小建议≥256

5. 实战技巧与性能优化

5.1 混合精度训练

scaler = torch.cuda.amp.GradScaler() with torch.cuda.amp.autocast(): outputs = model(inputs) loss = criterion(outputs, targets) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()

优势：

• 显存占用减少30-50%
• 训练速度提升20-40%
• 精度损失<0.5%

5.2 梯度检查点技术

from torch.utils.checkpoint import checkpoint def forward(self, x): x = checkpoint(self.blocks[:6], x) # 分段检查点 x = checkpoint(self.blocks[6:], x) return x

内存-计算权衡：

• 显存节省50-60%
• 训练时间增加约25%

5.3 不同硬件适配建议

6. 扩展与变体

6.1 分层ViT结构

class HierarchicalViT(nn.Module): def __init__(self): self.stage1 = PatchEmbed(56, 14) # 4x4 patches self.stage2 = PatchEmbed(28, 7) # 4x4->2x2 self.stage3 = PatchEmbed(14, 7) # 2x2->1x1

优势：

• 保留局部结构信息
• 计算量减少约40%
• 适合高分辨率图像

6.2 知识蒸馏方案

teacher = torch.hub.load('facebookresearch/deit:main', 'deit_base_patch16_224', pretrained=True) loss = 0.7 * cls_loss + 0.3 * distill_loss(teacher, student)

蒸馏策略对比：

• 软标签蒸馏：提升1-2%
• 注意力蒸馏：提升2-3%
• 隐藏层匹配：提升3-4%

在实际项目中，12层ViT配合适当的调参策略，在ImageNet上可以达到82%以上的top-1准确率。对于计算资源有限的场景，建议从6层模型开始逐步扩展，同时注意学习率与批量大小的协同调整。