深度学习之图像分类（十九）-- Vision Transformer(ViT)实战：从零构建与CIFAR-100数据集训练

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 tqdm2.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, trainTrue, downloadTrue, transformtransform) test_data datasets.CIFAR100(root./data, trainFalse, downloadTrue, transformtransform) train_loader DataLoader(train_data, batch_size64, shuffleTrue) test_loader DataLoader(test_data, batch_size64, shuffleFalse)注意实际项目中我会加入RandAugment等数据增强但为简化教程这里只使用基础变换。3. ViT模型实现3.1 Patch Embedding层这是将图像转换为token序列的关键步骤。以224x224输入为例使用16x16的patch大小会得到196个patchclass PatchEmbedding(nn.Module): def __init__(self, img_size224, patch_size16, in_channels3, embed_dim768): 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_sizepatch_size, stridepatch_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 x3.2 Transformer Encoder实现单个Encoder Block包含LayerNorm、Multi-Head Attention和MLPclass TransformerEncoder(nn.Module): def __init__(self, embed_dim768, num_heads12, mlp_ratio4.0, dropout0.1): super().__init__() self.norm1 nn.LayerNorm(embed_dim) self.attn nn.MultiheadAttention(embed_dim, num_heads, dropoutdropout) 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 x3.3 完整ViT模型整合所有组件并添加class token和position embeddingclass ViT(nn.Module): def __init__(self, img_size224, patch_size16, in_channels3, num_classes100, embed_dim768, depth12, num_heads12, mlp_ratio4.0, dropout0.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, std0.02) nn.init.trunc_normal_(self.cls_token, std0.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), dim1) 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_classes100).to(device) criterion nn.CrossEntropyLoss() optimizer optim.AdamW(model.parameters(), lr3e-4, weight_decay0.05) scheduler optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max100)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容易过拟合我总结了几个有效策略强数据增强加入RandAugment或MixUpDropPath随机丢弃部分网络路径权重衰减增大AdamW的weight_decay参数早停机制监控验证集loss# DropPath实现示例 def drop_path(x, drop_prob0., trainingFalse): 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, dtypex.dtype, devicex.device) random_tensor.floor_() output x.div(keep_prob) * random_tensor return output5. 结果分析与可视化训练完成后我们可以可视化注意力图来理解模型关注点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), dim1) 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, cmaphot) plt.colorbar() plt.show()在CIFAR-100上经过100个epoch训练后ViT-Base通常能达到约65-70%的测试准确率。虽然不如在ImageNet上的表现惊艳但这个结果已经超过了同等规模的CNN模型。