CIFAR-10图像分类实战：如何用Transformer模型超越CNN基准性能

1次阅读

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

在 CIFAR-10 这样的低分辨率（32×32）图像分类任务中，传统 CNN 通过堆叠卷积层逐渐扩大感受野。但实际存在三个明显瓶颈：

浅层网络难以捕获全局依赖关系（需 5 层以上卷积才能覆盖全图）
池化操作导致空间信息不可逆损失
平移不变性假设与细粒度分类需求存在矛盾

模型类型	参数量 (M)	FLOPs(G)	准确率 (%)
ResNet-18	11.2	0.56	94.2
ViT-Tiny	5.7	1.3	95.7
Swin-Tiny	6.9	1.8	96.1

上表基于相同训练配置（AdamW 优化器，300epochs）

class PatchEmbed(nn.Module):
    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)  # [B,192,8,8]
        self.norm = nn.LayerNorm(embed_dim)

    def forward(self, x):
        x = self.proj(x)
        x = x.flatten(2).transpose(1, 2)  # [B,64,192]
        return self.norm(x)

关键细节：
– 使用 4 ×4 步长卷积代替原始 ViT 的线性投影，保留局部结构信息
– LayerNorm 在投影后应用，比前置 Norm 更稳定

# 生成可学习的位置编码
pos_embed = nn.Parameter(torch.randn(1, 64+1, 192) * 0.02)

# 可视化前两个维度
plt.figure(figsize=(10,5))
plt.imshow(pos_embed[0,:,0:2].detach().numpy())
plt.colorbar()

可见分类 token（首行）与空间位置编码的明显差异

scaler = torch.cuda.amp.GradScaler()

for epoch in range(epochs):
    for x, y in train_loader:
        x, y = x.cuda(), y.cuda()

        with torch.cuda.amp.autocast():
            out = model(x)
            loss = criterion(out, y)

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