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

1次阅读

没有评论

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

CIFAR100 作为经典的图像分类数据集，包含 32×32 像素的小尺寸图像，这给模型设计带来了独特挑战：

空间信息有限 ：32×32 分辨率下，传统 CNN 的局部感受野难以捕获长距离依赖关系
类别间相似度高 ：100 个细粒度类别（如 20 个超类下的子类）导致类间差异细微
数据量瓶颈 ：50,000 张训练图像对 Transformer 这类数据饥渴型架构构成挑战

传统 CNN（如 ResNet）在 CIFAR100 上常出现：
1. 高层特征丢失细粒度信息
2. 跨区域特征关联建模能力不足
3. 对小物体旋转 / 尺度变化敏感

在 Transformer 变体中，我们选择标准 ViT 而非 Swin Transformer，主要考虑：

计算效率 ：Swin 的窗口注意力在 32×32 分辨率下优势不明显
实现复杂度 ：ViT 的全局注意力更适合学术验证
Patch 尺寸 ：实验表明 4 ×4 patches（共 64 个 token）在精度与速度间最佳平衡

关键计算量对比（FLOPs）：

模型	Patch 大小	FLOPs
ViT-Tiny	4×4	0.5G
Swin-Tiny	4×4	1.1G
ViT-Small	8×8	1.8G

import torch
import torch.nn as nn
from einops import rearrange

class PatchEmbedding(nn.Module):
    """将图像分割为 patch 并线性投影"""
    def __init__(self, img_size=32, patch_size=4, in_chans=3, embed_dim=192):
        super().__init__()
        self.proj = nn.Conv2d(in_chans, embed_dim, 
                             kernel_size=patch_size, 
                             stride=patch_size)
        self.num_patches = (img_size // patch_size) ** 2

    def forward(self, x):
        x = self.proj(x)  # [B, C, H, W] -> [B, E, H/P, W/P]
        x = rearrange(x, 'b e h w -> b (h w) e')
        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):
        # 残差连接 + 层归一化
        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.transforms import autoaugment, transforms

train_transform = transforms.Compose([transforms.RandomCrop(32, padding=4, padding_mode='reflect'),
    transforms.RandomHorizontalFlip(),
    autoaugment.RandAugment(
        num_ops=2,  # 操作次数减少
        magnitude=9,  # 强度降低
        num_magnitude_bins=31,
        interpolation=transforms.InterpolationMode.BILINEAR
    ),
    transforms.ToTensor(),
    transforms.Normalize((0.5071, 0.4867, 0.4408), (0.2675, 0.2565, 0.2761))
])

class LabelSmoothingCrossEntropy(nn.Module):
    def __init__(self, epsilon=0.1):
        super().__init__()
        self.epsilon = epsilon

    def forward(self, logits, targets):
        log_probs = -torch.log_softmax(logits, dim=-1)
        nll_loss = log_probs.gather(dim=-1, index=targets.unsqueeze(1))
        smooth_loss = log_probs.mean(dim=-1)
        loss = (1 - self.epsilon) * nll_loss + self.epsilon * smooth_loss
        return loss.mean()

scaler = torch.cuda.amp.GradScaler()

for inputs, targets in train_loader:
    optimizer.zero_grad()

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

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

effective_batch_size = 512
accum_steps = effective_batch_size // batch_size

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

    scaler.scale(loss).backward()

    if (i + 1) % accum_steps == 0:
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad()

学习率 warmup：
错误做法：直接从峰值学习率开始

正确实现：线性 warmup 5-10 个 epoch

scheduler = torch.optim.lr_scheduler.SequentialLR(
    optimizer,
    schedulers=[torch.optim.lr_scheduler.LinearLR(optimizer, 0.001, 1, total_iters=5),
        torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs-5)
    ],
    milestones=[5]
)