# Swin Transformer实战：如何用局部自注意力提升图像分类效果（附代码） 在计算机视觉领域，Transformer架构正逐渐取代传统CNN成为主流选择。然而，全局自注意力机制的高计算复杂度一直是制约其发展的瓶颈。本文将深入探讨Swin Transformer中创新的局部自注意力机制，并通过完整代码示例展示如何在实际图像分类任务中实现性能突破。 ## 1. 局部自注意力机制的核心原理 局部自注意力（Local Self-Attention）通过将输入划分为多个不重叠窗口，在每个窗口内独立计算注意力权重。这种设计带来了两大优势： - **计算效率**：将O(n²)复杂度降为O(n×w²)，其中w为窗口大小 - **局部特征捕捉**：更聚焦于邻近像素关系，适合视觉任务的局部性特点 与传统CNN的固定权重卷积不同，局部自注意力具有动态权重调整能力。下表对比了三种主要特征提取方式： | 特性 | 标准卷积 | 全局自注意力 | 局部自注意力 | |---------------------|----------|--------------|--------------| | 感受野范围 | 固定 | 全局 | 局部可变 | | 权重动态性 | 静态 | 完全动态 | 局部动态 | | 计算复杂度 | O(k²×n) | O(n²) | O(n×w²) | | 内存占用 | 低 | 极高 | 中等 | ```python # 基础局部自注意力实现 import torch import torch.nn as nn class LocalSelfAttention(nn.Module): def __init__(self, dim, window_size=7): super().__init__() self.window_size = window_size self.qkv = nn.Linear(dim, dim*3) self.proj = nn.Linear(dim, dim) def forward(self, x): B, H, W, C = x.shape x = x.view(B, H//self.window_size, self.window_size, W//self.window_size, self.window_size, C) x = x.permute(0,1,3,2,4,5).reshape(-1, self.window_size*self.window_size, C) qkv = self.qkv(x).chunk(3, dim=-1) q, k, v = map(lambda t: t.view(-1, self.window_size**2, C), qkv) attn = (q @ k.transpose(-2,-1)) * (C**-0.5) attn = attn.softmax(dim=-1) out = (attn @ v).view(B, H, W, C) return self.proj(out) ``` > 提示：窗口大小的选择需要平衡计算效率和特征捕捉能力。常见图像分类任务中，7×7或8×8窗口表现最佳 ## 2. Swin Transformer的滑动窗口创新 Swin Transformer通过两种关键技术解决了普通局部自注意力的局限性： **2.1 层级特征金字塔** 构建了4个阶段的特征图，每阶段下采样2倍，形成典型的多尺度表示： ``` Stage 1: 56×56特征图 Stage 2: 28×28特征图 Stage 3: 14×14特征图 Stage 4: 7×7特征图 ``` **2.2 滑动窗口机制** 在连续Transformer块中交替使用两种窗口划分方式： - 常规窗口划分：均匀分割特征图 - 偏移窗口划分：窗口向右下角偏移50%窗口大小 这种设计实现了跨窗口信息交互，同时保持计算效率。下图展示了两种窗口模式的交替： ``` 常规窗口： 偏移窗口： ┌───┬───┐ ┌─┬─┬─┐ │ 1 │ 2 │ ├─┼─┼─┤ ├───┼───┤ │4│5│6│ │ 3 │ 4 │ ├─┼─┼─┤ └───┴───┘ └─┴─┴─┘ ``` ```python # 滑动窗口实现关键代码 def create_mask(H, W, window_size, shift_size): img_mask = torch.zeros((1, H, W, 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. 完整图像分类实现 以下是在ImageNet数据集上实现Swin-Tiny模型的完整流程： **3.1 模型架构** ```python class SwinTransformer(nn.Module): def __init__(self, img_size=224, patch_size=4, in_chans=3, num_classes=1000, embed_dim=96, depths=[2,2,6,2], num_heads=[3,6,12,24], window_size=7, mlp_ratio=4.): super().__init__() self.patch_embed = PatchEmbed(img_size, patch_size, in_chans, embed_dim) self.pos_drop = nn.Dropout(p=0.1) self.layers = nn.ModuleList() for i in range(4): layer = BasicLayer( dim=int(embed_dim * 2**i), depth=depths[i], num_heads=num_heads[i], window_size=window_size, mlp_ratio=mlp_ratio, drop_path_rate=0.1 ) self.layers.append(layer) self.norm = nn.LayerNorm(embed_dim*8) self.avgpool = nn.AdaptiveAvgPool1d(1) self.head = nn.Linear(embed_dim*8, num_classes) def forward(self, x): x = self.patch_embed(x) x = self.pos_drop(x) for layer in self.layers: x = layer(x) x = self.norm(x) x = self.avgpool(x.transpose(1,2)) x = torch.flatten(x, 1) x = self.head(x) return x ``` **3.2 关键训练参数配置** ```yaml # configs/swin_tiny.yaml data: train_root: /path/to/imagenet/train val_root: /path/to/imagenet/val batch_size: 256 num_workers: 8 model: img_size: 224 patch_size: 4 embed_dim: 96 depths: [2, 2, 6, 2] num_heads: [3, 6, 12, 24] training: epochs: 300 lr: 5e-4 weight_decay: 0.05 label_smoothing: 0.1 mixup: 0.8 cutmix: 1.0 ``` **3.3 数据增强策略** - RandAugment：随机选择2-3种增强操作 - MixUp & CutMix：混合样本增强 - Random Erasing：随机区域擦除 - Color Jitter：颜色扰动 ```python from torchvision import transforms train_transform = transforms.Compose([ transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), transforms.RandAugment(), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), transforms.RandomErasing(p=0.2) ]) ``` ## 4. 性能优化技巧 **4.1 混合精度训练** ```python 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() ``` **4.2 梯度裁剪** ```python torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) ``` **4.3 学习率调度** 采用余弦退火配合线性warmup： ```python from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR warmup_epochs = 5 scheduler1 = LinearLR(optimizer, start_factor=0.01, total_iters=warmup_epochs) scheduler2 = CosineAnnealingLR(optimizer, T_max=epochs-warmup_epochs) scheduler = SequentialLR(optimizer, [scheduler1, scheduler2], [warmup_epochs]) ``` **4.4 模型量化推理** ```python quantized_model = torch.quantization.quantize_dynamic( model, {torch.nn.Linear}, dtype=torch.qint8 ) ``` 在实际测试中，Swin-Tiny模型在ImageNet-1K上达到81.3% top-1准确率，相比ResNet-50提升约3%，而计算量仅增加15%。这种性能提升在细粒度分类任务中更为显著，如CUB-200数据集上可提升5-7个百分点。