告别CNN？手把手带你用PyTorch复现ViT（Vision Transformer）图像分类模型

从零构建ViT模型：PyTorch实战图像分类新范式

当你在Instagram上传照片时，那个能自动识别出猫、狗或风景的AI系统，很可能基于卷积神经网络(CNN)。但今天，我们要挑战这个持续了三十年的视觉处理范式。2017年Transformer在NLP领域的爆发，终于在2020年通过Vision Transformer(ViT)彻底改写了图像处理的游戏规则。

1. 环境准备与数据预处理

在开始构建ViT之前，确保你的开发环境已安装PyTorch 1.8+和TorchVision。对于GPU加速，建议使用CUDA 11.x：

conda create -n vit python=3.8 conda install pytorch torchvision cudatoolkit=11.3 -c pytorch pip install einops matplotlib tqdm

我们将使用CIFAR-10数据集作为示例，这个经典数据集包含60,000张32x32像素的彩色图像，分为10个类别。与ImageNet相比，它体积小但足够验证模型有效性：

from torchvision import datasets, transforms train_transform = transforms.Compose([ transforms.RandomHorizontalFlip(), transforms.RandomRotation(15), transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) ]) train_set = datasets.CIFAR10(root='./data', train=True, download=True, transform=train_transform) test_set = datasets.CIFAR10(root='./data', train=False, download=True, transform=train_transform)

ViT与传统CNN最大的预处理差异在于图像分块(Patch Embedding)。对于32x32的CIFAR-10图像，如果我们选择8x8的patch大小，将得到16个patch(32/8=4，4x4=16)：

import torch import torch.nn as nn class PatchEmbedding(nn.Module): def __init__(self, img_size=32, patch_size=8, in_channels=3, embed_dim=128): super().__init__() self.img_size = img_size self.patch_size = patch_size self.n_patches = (img_size // patch_size) ** 2 self.proj = nn.Conv2d( in_channels, embed_dim, kernel_size=patch_size, stride=patch_size ) def forward(self, x): x = self.proj(x) # (B, E, H/P, W/P) x = x.flatten(2) # (B, E, N) x = x.transpose(1, 2) # (B, N, E) return x

提示：patch_size的选择需要权衡模型性能和计算复杂度。较小的patch能保留更多细节但会增加序列长度，通常建议在8-16像素之间选择。

2. ViT核心组件实现

2.1 位置编码的创新实现

Transformer原本是为序列数据设计的，缺乏对2D图像结构的理解。ViT通过位置编码(position embedding)来解决这个问题。不同于原始Transformer的1D位置编码，我们实现了更适应图像的方案：

class PositionalEncoding(nn.Module): def __init__(self, n_patches=16, embed_dim=128): super().__init__() self.pos_embed = nn.Parameter(torch.zeros(1, n_patches + 1, embed_dim)) nn.init.trunc_normal_(self.pos_embed, std=0.02) def forward(self, x): return x + self.pos_embed[:, :x.size(1)]

实际应用中，我们发现几种位置编码变体的效果对比：

2.2 Transformer编码器详解

ViT的核心是由多个Transformer Encoder层堆叠而成。每个Encoder包含多头注意力(MHA)和前馈网络(FFN)：

class TransformerBlock(nn.Module): def __init__(self, embed_dim=128, num_heads=4, mlp_ratio=4.0, dropout=0.1): super().__init__() self.norm1 = nn.LayerNorm(embed_dim) self.attn = nn.MultiheadAttention(embed_dim, num_heads, dropout=dropout) self.norm2 = nn.LayerNorm(embed_dim) self.mlp = nn.Sequential( nn.Linear(embed_dim, int(embed_dim * mlp_ratio)), nn.GELU(), nn.Dropout(dropout), nn.Linear(int(embed_dim * mlp_ratio), embed_dim), nn.Dropout(dropout) ) def forward(self, x): res = x x = self.norm1(x) x, _ = self.attn(x, x, x) x = res + x res = x x = self.norm2(x) x = self.mlp(x) x = res + x return x

关键参数配置建议：

• embed_dim: 128-512 (根据可用GPU内存调整)
• num_heads: 4-12 (通常选择embed_dim能被整除的值)
• mlp_ratio: 2.0-4.0 (控制FFN层的扩展倍数)
• depth: 6-12层 (更深的网络需要更多数据)

3. 完整ViT模型组装

现在我们将各个组件组合成完整模型，并添加分类头：

class VisionTransformer(nn.Module): def __init__(self, img_size=32, patch_size=8, in_channels=3, embed_dim=128, depth=6, num_heads=4, mlp_ratio=4.0, num_classes=10, dropout=0.1): super().__init__() self.patch_embed = PatchEmbedding(img_size, patch_size, in_channels, embed_dim) self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) self.pos_embed = PositionalEncoding(self.patch_embed.n_patches, embed_dim) self.blocks = nn.ModuleList([ TransformerBlock(embed_dim, num_heads, mlp_ratio, dropout) 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) # (B, N, E) cls_token = self.cls_token.expand(x.size(0), -1, -1) x = torch.cat((cls_token, x), dim=1) # (B, 1+N, E) x = self.pos_embed(x) for block in self.blocks: x = block(x) x = self.norm(x) cls_token_final = x[:, 0] # 取出分类token return self.head(cls_token_final)

注意：cls_token是ViT的关键设计之一，它作为一个可学习的参数，通过自注意力机制聚合全局信息，最终用于分类决策。

4. 训练策略与调优技巧

4.1 优化器配置与学习率调度

ViT对优化策略非常敏感。我们推荐使用AdamW优化器配合余弦退火学习率调度：

from torch.optim import AdamW from torch.optim.lr_scheduler import CosineAnnealingLR model = VisionTransformer().to(device) optimizer = AdamW(model.parameters(), lr=3e-4, weight_decay=0.05) scheduler = CosineAnnealingLR(optimizer, T_max=100, eta_min=1e-5) criterion = nn.CrossEntropyLoss()

实验表明，不同的优化配置对最终准确率影响显著：

4.2 数据增强与正则化

由于ViT缺乏CNN固有的平移不变性等归纳偏置，数据增强尤为重要：

from torchvision import transforms train_transform = transforms.Compose([ transforms.RandomResizedCrop(32, scale=(0.8, 1.0)), transforms.RandomHorizontalFlip(), transforms.RandomRotation(15), transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2), transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), transforms.RandomErasing(p=0.1) ])

在CIFAR-10上，不同正则化技术的效果对比：

1. Dropout：在注意力层和FFN层后添加，通常设为0.1
2. Stochastic Depth：随机跳过某些层，缓解过拟合
3. Layer Scale：对残差连接进行缩放，稳定深层训练
4. MixUp：图像混合增强，提升模型鲁棒性

4.3 梯度裁剪与混合精度训练

ViT训练过程中容易出现梯度爆炸，梯度裁剪至关重要：

torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

同时，使用混合精度训练可以大幅减少显存占用并加速训练：

from torch.cuda.amp import GradScaler, autocast scaler = GradScaler() for inputs, labels in train_loader: inputs, labels = inputs.to(device), labels.to(device) with autocast(): outputs = model(inputs) loss = criterion(outputs, labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update() optimizer.zero_grad()

5. 模型评估与结果分析

在CIFAR-10测试集上评估我们实现的ViT模型：

model.eval() correct = 0 total = 0 with torch.no_grad(): for images, labels in test_loader: images, labels = images.to(device), labels.to(device) outputs = model(images) _, predicted = torch.max(outputs.data, 1) total += labels.size(0) correct += (predicted == labels).sum().item() print(f'Test Accuracy: {100 * correct / total:.2f}%')

与常见模型的对比结果：

可视化注意力图可以帮助我们理解模型关注的重点区域：

def visualize_attention(model, image): model.eval() with torch.no_grad(): # 获取最后一层的注意力权重 attn_weights = model.blocks[-1].attn.get_attention_map(image.unsqueeze(0)) # 将注意力权重映射回图像空间 patch_size = model.patch_embed.patch_size heatmap = attn_weights[0, 0, 1:].reshape(4, 4).cpu().numpy() heatmap = cv2.resize(heatmap, (32, 32)) plt.imshow(image.permute(1, 2, 0).cpu().numpy()) plt.imshow(heatmap, alpha=0.5, cmap='jet') plt.show()

在实际项目中，我们发现ViT在以下场景表现尤为突出：

• 需要全局上下文理解的任务（如场景分类）
• 数据量充足的情况下（>1M图像）
• 对模型可解释性要求较高的应用

而在以下场景CNN可能仍是更好选择：

• 数据量有限（<100K图像）
• 需要实时推理的移动端应用
• 对局部纹理特征敏感的任务（如细粒度分类）