1. ViT模型实战概述

第一次接触Vision Transformer时，我被它彻底抛弃卷积操作的勇气震惊了。这个将NLP领域的Transformer直接搬到CV领域的"暴力美学"方案，在ImageNet上竟然能达到88.55%的top-1准确率。今天我们就用PyTorch从零实现ViT，并在CIFAR-100数据集上完成训练实战。

为什么选择CIFAR-100？这个包含100个类别的数据集虽然分辨率只有32x32，但正因如此，它特别适合验证ViT在小规模数据上的表现。我在实际项目中多次使用这个组合测试模型性能，发现即使没有海量预训练数据，通过合理的调参ViT也能展现出惊人的潜力。

2. 环境准备与数据加载

2.1 安装依赖库

建议使用Python 3.8+和PyTorch 1.10+环境。以下是必须的依赖项：

pip install torch torchvision matplotlib numpy tqdm

我习惯在代码开头统一导入所需模块：

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import numpy as np
import matplotlib.pyplot as plt
from tqdm import tqdm

2.2 CIFAR-100数据处理

ViT通常接受224x224的输入，但CIFAR-100原始尺寸为32x32。经过多次实验，我发现直接上采样效果反而比保持原尺寸更好：

transform = transforms.Compose([
    transforms.Resize(224),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

train_data = datasets.CIFAR100(root='./data', train=True, download=True, transform=transform)
test_data = datasets.CIFAR100(root='./data', train=False, download=True, transform=transform)

train_loader = DataLoader(train_data, batch_size=64, shuffle=True)
test_loader = DataLoader(test_data, batch_size=64, shuffle=False)

注意：实际项目中我会加入RandAugment等数据增强，但为简化教程这里只使用基础变换。

3. ViT模型实现

3.1 Patch Embedding层

这是将图像转换为token序列的关键步骤。以224x224输入为例，使用16x16的patch大小会得到196个patch：

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, img_size)
        self.patch_size = (patch_size, patch_size)
        self.num_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):
        B, C, H, W = x.shape
        assert H == self.img_size[0] and W == self.img_size[1], \
            f"输入图像尺寸({H}*{W})与模型设定({self.img_size[0]}*{self.img_size[1]})不符"
        
        x = self.proj(x).flatten(2).transpose(1, 2)  # [B, C, H, W] -> [B, num_patches, embed_dim]
        return x

3.2 Transformer Encoder实现

单个Encoder Block包含LayerNorm、Multi-Head Attention和MLP：

class TransformerEncoder(nn.Module):
    def __init__(self, embed_dim=768, num_heads=12, 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)
        )
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, x):
        x = x + self.dropout(self.attn(self.norm1(x), self.norm1(x), self.norm1(x))[0])
        x = x + self.dropout(self.mlp(self.norm2(x)))
        return x

3.3 完整ViT模型

整合所有组件并添加class token和position embedding：

class ViT(nn.Module):
    def __init__(self, img_size=224, patch_size=16, in_channels=3, num_classes=100,
                 embed_dim=768, depth=12, num_heads=12, mlp_ratio=4.0, 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 = nn.Parameter(torch.zeros(1, self.patch_embed.num_patches + 1, embed_dim))
        self.pos_drop = nn.Dropout(dropout)
        
        self.blocks = nn.ModuleList([
            TransformerEncoder(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)
        
        nn.init.trunc_normal_(self.pos_embed, std=0.02)
        nn.init.trunc_normal_(self.cls_token, std=0.02)
    
    def forward(self, x):
        B = x.shape[0]
        x = self.patch_embed(x)
        
        cls_tokens = self.cls_token.expand(B, -1, -1)
        x = torch.cat((cls_tokens, x), dim=1)
        x = x + self.pos_embed
        x = self.pos_drop(x)
        
        for blk in self.blocks:
            x = blk(x)
        
        x = self.norm(x)
        return self.head(x[:, 0])

4. 模型训练与调优

4.1 训练配置

使用AdamW优化器和余弦学习率衰减：

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = ViT(num_classes=100).to(device)

criterion = nn.CrossEntropyLoss()
optimizer = optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.05)

scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100)

4.2 训练循环

我习惯在每个epoch后验证模型表现：

def train_epoch(model, loader, optimizer, criterion, device):
    model.train()
    total_loss = 0
    correct = 0
    
    for images, labels in tqdm(loader):
        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()
        _, predicted = outputs.max(1)
        correct += predicted.eq(labels).sum().item()
    
    return total_loss / len(loader), correct / len(loader.dataset)

4.3 解决过拟合问题

在小数据集上训练ViT容易过拟合，我总结了几个有效策略：

1. 强数据增强 ：加入RandAugment或MixUp
2. DropPath ：随机丢弃部分网络路径
3. 权重衰减 ：增大AdamW的weight_decay参数
4. 早停机制 ：监控验证集loss

# DropPath实现示例
def drop_path(x, drop_prob=0., training=False):
    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

5. 结果分析与可视化

训练完成后，我们可以可视化注意力图来理解模型关注点：

def visualize_attention(model, image):
    model.eval()
    with torch.no_grad():
        feature = model.patch_embed(image.unsqueeze(0))
        cls_token = model.cls_token.expand(1, -1, -1)
        x = torch.cat((cls_token, feature), dim=1)
        x = x + model.pos_embed
        x = model.pos_drop(x)
        
        attentions = []
        for blk in model.blocks:
            x = blk.norm1(x)
            _, attn = blk.attn(x, x, x)
            attentions.append(attn)
    
    # 可视化最后一个block的注意力
    attn = attentions[-1][0, :, 0, 1:].mean(0)  # 取cls token对其他patch的注意力
    attn = attn.reshape(14, 14).cpu().numpy()
    
    plt.imshow(attn, cmap='hot')
    plt.colorbar()
    plt.show()

在CIFAR-100上，经过100个epoch训练后，ViT-Base通常能达到约65-70%的测试准确率。虽然不如在ImageNet上的表现惊艳，但这个结果已经超过了同等规模的CNN模型。