基于Transformer的CIFAR100图像分类实战：从模型选型到性能优化

1次阅读

共计 2990 个字符，预计需要花费 8 分钟才能阅读完成。

CIFAR100 作为经典的图像分类数据集，包含 32×32 像素的小尺寸图像，共 100 个类别。这类低分辨率图像分类任务面临两个核心挑战：

局部信息模糊 ：32×32 的输入尺寸导致单个像素承载的语义信息有限，传统 CNN 的局部感受野难以捕获有效特征
长距离依赖缺失 ：CNN 通过堆叠卷积层建立远程关系，但深层网络容易引发梯度消失，且计算成本呈平方增长

相比而言，Transformer 的 Self-Attention 机制能直接建模全局依赖关系，通过可学习的注意力权重显式建立像素间关联。实验表明，在相同计算预算下，ViT 模型的注意力头能覆盖整个图像空间，而 CNN 的有效感受野仅能覆盖局部区域。

在多种视觉 Transformer 变体中，选择标准 ViT 架构基于以下考虑：

计算效率 ：Swin Transformer 的窗口划分在 32×32 分辨率下收益有限，而 ViT 的 16×16 patch size（共 4 个 patch）在 CIFAR100 上已达合理粒度
实现简洁 ：原生 ViT 的线性投影层比 Swin 的移位窗口更易于调试，适合快速验证基线性能
硬件适配 ：小 patch 数量使得模型在消费级 GPU（如 RTX 3090）上即可完成训练，无需分布式训练

关键计算公式：

 计算复杂度 = O(n²·d)  # n 为序列长度，d 为特征维度
当 patch_size=16 时，n=4（32/16）²，显著低于 224x224 输入时的 196

import torch
import torch.nn as nn

class PatchEmbedding(nn.Module):
    """将图像分割为 patch 并线性投影"""
    def __init__(self, img_size=32, patch_size=16, in_chans=3, embed_dim=192):
        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, D, H/p, W/p]
        x = x.flatten(2).transpose(1, 2)  # [B, D, N] -> [B, N, D]
        return x

class TransformerBlock(nn.Module):
    """标准 Transformer 编码器块"""
    def __init__(self, dim, num_heads, mlp_ratio=4., drop=0.):
        super().__init__()
        self.norm1 = nn.LayerNorm(dim)
        self.attn = nn.MultiheadAttention(dim, num_heads, dropout=drop)
        self.norm2 = nn.LayerNorm(dim)
        self.mlp = nn.Sequential(nn.Linear(dim, int(dim * mlp_ratio)),
            nn.GELU(),
            nn.Dropout(drop),
            nn.Linear(int(dim * mlp_ratio), dim)
        )

    def forward(self, x):
        # Pre-Norm 结构
        x = x + self.attn(self.norm1(x), self.norm1(x), self.norm1(x))[0]
        x = x + self.mlp(self.norm2(x))
        return x

针对小尺寸图像优化的 RandAugment 配置：

from torchvision import transforms

train_transform = transforms.Compose([transforms.RandomCrop(32, padding=4),
    transforms.RandomHorizontalFlip(),
    transforms.RandAugment(
        num_ops=2,  # 操作次数减少以避免过度扰动
        magnitude=9  # 强度降低适应小图像
    ),
    transforms.ToTensor(),
    transforms.Normalize((0.5071, 0.4867, 0.4408), (0.2675, 0.2565, 0.2761))
])

class LabelSmoothingCrossEntropy(nn.Module):
    def __init__(self, smoothing=0.1):
        super().__init__()
        self.confidence = 1.0 - smoothing
        self.smoothing = smoothing

    def forward(self, x, target):
        logprobs = F.log_softmax(x, dim=-1)
        nll_loss = -logprobs.gather(dim=-1, index=target.unsqueeze(1))
        nll_loss = nll_loss.squeeze(1)
        smooth_loss = -logprobs.mean(dim=-1)
        loss = self.confidence * nll_loss + self.smoothing * smooth_loss
        return loss.mean()

scaler = torch.cuda.amp.GradScaler()

for inputs, targets in train_loader:
    with torch.cuda.amp.autocast():
        outputs = model(inputs)
        loss = criterion(outputs, targets)

    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

当 GPU 显存不足时，通过多次前向传播累积梯度：

gradient_accum_steps = 4

for i, (inputs, targets) in enumerate(train_loader):
    loss = model(inputs, targets) / gradient_accum_steps
    loss.backward()

    if (i+1) % gradient_accum_steps == 0:
        optimizer.step()
        optimizer.zero_grad()