从零到一：Swin Transformer图像分类实战（PyTorch版）

从零到一：Swin Transformer图像分类实战（PyTorch版）

在计算机视觉领域，Transformer架构正逐渐取代传统的CNN成为新的主流。Swin Transformer作为微软亚洲研究院提出的创新模型，通过分层特征映射和移位窗口机制，在图像分类任务中展现出卓越性能。本文将带您从零开始，完整实现基于PyTorch的Swin Transformer图像分类解决方案。

1. 环境配置与准备工作

首先需要搭建适合深度学习开发的环境。推荐使用Anaconda创建独立的Python环境，避免依赖冲突：

conda create -n swin python=3.8 conda activate swin

安装核心依赖包时，特别注意PyTorch与CUDA版本的匹配：

pip install torch==1.10.0+cu113 torchvision==0.11.1+cu113 -f https://download.pytorch.org/whl/torch_stable.html pip install timm matplotlib opencv-python

硬件配置建议：

• GPU：NVIDIA RTX 3060及以上（显存≥8GB）
• 内存：16GB以上
• 存储：SSD硬盘，至少50GB可用空间

提示：可使用nvidia-smi命令检查GPU状态，确保CUDA驱动正常

2. 数据准备与预处理

高质量的数据准备是模型成功的关键。我们以花卉分类数据集为例，展示标准处理流程：

from torchvision import transforms train_transform = transforms.Compose([ transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) val_transform = transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ])

数据集组织结构应遵循以下规范：

data/flower_photos/ ├── daisy/ ├── dandelion/ ├── roses/ ├── sunflowers/ └── tulips/

数据加载器实现示例：

from torch.utils.data import DataLoader train_loader = DataLoader( train_dataset, batch_size=32, shuffle=True, num_workers=4, pin_memory=True ) val_loader = DataLoader( val_dataset, batch_size=32, shuffle=False, num_workers=4 )

3. 模型构建与核心原理

Swin Transformer的核心创新在于其分层设计和移位窗口机制。让我们深入解析模型的关键组件：

3.1 基础模块实现

窗口多头注意力(Window Multi-Head Attention)实现：

class WindowAttention(nn.Module): def __init__(self, dim, window_size, num_heads): super().__init__() self.dim = dim self.window_size = window_size self.num_heads = num_heads self.scale = (dim // num_heads) ** -0.5 self.qkv = nn.Linear(dim, dim * 3) self.proj = nn.Linear(dim, dim) # 相对位置编码 self.relative_position_bias_table = nn.Parameter( torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads)) # 生成相对位置索引 coords = torch.stack(torch.meshgrid( [torch.arange(window_size[0]), torch.arange(window_size[1])])) coords_flatten = torch.flatten(coords, 1) relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] relative_coords = relative_coords.permute(1, 2, 0).contiguous() relative_coords[:, :, 0] += window_size[0] - 1 relative_coords[:, :, 1] += window_size[1] - 1 relative_coords[:, :, 0] *= 2 * window_size[1] - 1 relative_position_index = relative_coords.sum(-1) self.register_buffer("relative_position_index", relative_position_index)

移位窗口分区函数：

def create_mask(H, W, window_size, shift_size): Hp = int(np.ceil(H / window_size)) * window_size Wp = int(np.ceil(W / window_size)) * window_size img_mask = torch.zeros((1, Hp, Wp, 1)) h_slices = (slice(0, -window_size), slice(-window_size, -shift_size), slice(-shift_size, None)) w_slices = (slice(0, -window_size), slice(-window_size, -shift_size), slice(-shift_size, None)) cnt = 0 for h in h_slices: for w in w_slices: img_mask[:, h, w, :] = cnt cnt += 1 mask_windows = window_partition(img_mask, window_size) mask_windows = mask_windows.view(-1, window_size * window_size) attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2) attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)) return attn_mask

3.2 完整模型架构

Swin Transformer采用分层金字塔结构，各阶段配置如下表所示：

模型初始化代码：

def swin_tiny_patch4_window7_224(num_classes=1000): model = SwinTransformer( patch_size=4, in_chans=3, num_classes=num_classes, embed_dim=96, depths=[2, 2, 6, 2], num_heads=[3, 6, 12, 24], window_size=7, mlp_ratio=4., qkv_bias=True, drop_rate=0.0, drop_path_rate=0.1 ) return model

4. 模型训练与优化

4.1 训练策略配置

采用AdamW优化器配合余弦退火学习率调度：

optimizer = torch.optim.AdamW( model.parameters(), lr=5e-4, weight_decay=0.05 ) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max=20, eta_min=1e-6 )

关键训练参数：

• Batch Size: 32-128（根据GPU显存调整）
• Epochs: 50-100
• 混合精度训练：使用torch.cuda.amp加速

4.2 训练过程监控

使用TensorBoard记录训练指标：

from torch.utils.tensorboard import SummaryWriter writer = SummaryWriter() for epoch in range(epochs): train_loss, train_acc = train_one_epoch(model, train_loader, optimizer) val_loss, val_acc = validate(model, val_loader) writer.add_scalar('Loss/train', train_loss, epoch) writer.add_scalar('Accuracy/train', train_acc, epoch) writer.add_scalar('Loss/val', val_loss, epoch) writer.add_scalar('Accuracy/val', val_acc, epoch) scheduler.step()

4.3 模型评估指标

除了常规的准确率，建议监控以下指标：

5. 模型部署与推理

训练完成后，可以使用以下代码进行单张图片预测：

def predict(image_path, model, transform): img = Image.open(image_path).convert('RGB') img_tensor = transform(img).unsqueeze(0) with torch.no_grad(): output = model(img_tensor) probs = torch.nn.functional.softmax(output, dim=1) return probs.numpy()

部署优化技巧：

• 使用torch.jit.trace进行模型序列化
• 启用torch.inference_mode()提升推理速度
• 对输入进行批处理(batch inference)提高吞吐量

实际项目中，模型部署的典型性能指标：

6. 进阶优化技巧

6.1 数据增强策略

from timm.data.auto_augment import rand_augment_transform rand_augment = rand_augment_transform( config_str='rand-m9-mstd0.5', hparams={'img_mean': (0.485, 0.456, 0.406)} ) train_transform.transforms.insert(0, rand_augment)

6.2 知识蒸馏应用

teacher_model = swin_base_patch4_window7_224(pretrained=True) student_model = swin_tiny_patch4_window7_224() distill_loss = nn.KLDivLoss(reduction='batchmean') def compute_distill_loss(teacher_logits, student_logits, temperature=3.0): soft_teacher = F.softmax(teacher_logits / temperature, dim=1) soft_student = F.log_softmax(student_logits / temperature, dim=1) return distill_loss(soft_student, soft_teacher) * (temperature ** 2)

6.3 混合精度训练

from torch.cuda.amp import autocast, GradScaler scaler = GradScaler() for inputs, targets in train_loader: optimizer.zero_grad() with autocast(): outputs = model(inputs) loss = criterion(outputs, targets) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()

7. 常见问题排查

训练过程中可能遇到的问题及解决方案：

1. Loss不下降

   • 检查学习率是否合适
   • 验证数据预处理是否正确
   • 确认模型参数是否正常更新

2. GPU显存不足

   • 减小batch size
   • 使用梯度累积
   • 尝试混合精度训练

3. 验证集性能波动大

   • 增加验证集样本量
   • 检查数据泄露问题
   • 调整正则化强度

注意：当出现NaN值时，应立即中断训练，检查数据范围和模型结构

以下是一个典型训练过程的超参数配置参考：

config = { 'batch_size': 64, 'lr': 5e-4, 'weight_decay': 0.05, 'epochs': 100, 'warmup_epochs': 5, 'min_lr': 1e-6, 'clip_grad': 5.0, 'drop_path_rate': 0.2 }

通过本教程，您应该已经掌握了Swin Transformer在图像分类任务中的完整实现流程。在实际应用中，可根据具体任务需求调整模型结构和训练策略。