别再纠结选CNN还是Transformer了！手把手带你用PyTorch复现CoAtNet，感受‘混合双打’的魅力

从理论到实践：用PyTorch构建CoAtNet混合架构的完整指南

在计算机视觉领域，卷积神经网络(CNN)和Transformer架构长期处于竞争状态。CNN凭借其局部感受野和平移等变性在图像处理中表现出色，而Transformer则通过自注意力机制捕获长距离依赖关系。2021年出现的CoAtNet创新性地将两者优势结合，在ImageNet等基准测试中刷新了记录。本文将带您从零开始实现这个混合架构，通过代码级解析揭示其设计精髓。

1. 混合架构设计原理

1.1 CNN与Transformer的互补优势

传统CNN通过卷积核的滑动窗口操作获取局部特征，其固有特性包括：

• 平移等变性：物体位置变化时，特征响应同步变化
• 局部感知：3×3或5×5的有限感受野
• 参数共享：相同卷积核应用于所有位置

相比之下，Vision Transformer的核心优势在于：

• 全局上下文：自注意力机制建立任意两个图像块间的联系
• 动态权重：注意力分数根据输入内容自适应计算
• 并行处理：无需顺序处理特征图

# 传统卷积操作示例 import torch.nn as nn conv = nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, stride=1, padding=1) # 自注意力简化实现 def self_attention(Q, K, V): scores = torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(d_k) attn = torch.softmax(scores, dim=-1) return torch.matmul(attn, V)

1.2 MBConv块解析

CoAtNet采用MobileNetV2提出的MBConv作为基础模块，其结构包含：

1. 1×1扩展卷积（通常扩展4倍）
2. 3×3深度可分离卷积
3. SE(Squeeze-Excitation)注意力模块
4. 1×1投影卷积

class MBConv(nn.Module): def __init__(self, in_channels, out_channels, expansion=4): super().__init__() hidden_dim = in_channels * expansion self.block = nn.Sequential( nn.Conv2d(in_channels, hidden_dim, 1), nn.BatchNorm2d(hidden_dim), nn.SiLU(), nn.Conv2d(hidden_dim, hidden_dim, 3, padding=1, groups=hidden_dim), nn.BatchNorm2d(hidden_dim), nn.SiLU(), SqueezeExcite(hidden_dim), nn.Conv2d(hidden_dim, out_channels, 1), nn.BatchNorm2d(out_channels) )

提示：MBConv中的深度可分离卷积将标准卷积分解为depthwise和pointwise两步，计算量减少为原来的1/8到1/9

1.3 相对自注意力改进

标准Transformer中的绝对位置编码在图像任务中存在局限，CoAtNet采用相对位置编码：

$$ Attention = Softmax(\frac{QK^T + B}{\sqrt{d_k}})V $$

其中B是基于相对位置的偏置矩阵，让模型更好地理解像素间的空间关系。

2. PyTorch实现详解

2.1 整体架构搭建

CoAtNet采用分阶段设计，各阶段配置如下表：

class CoAtNet(nn.Module): def __init__(self, image_size=224): super().__init__() self.s0 = nn.Sequential( nn.Conv2d(3, 64, 3, stride=2, padding=1), nn.BatchNorm2d(64), nn.SiLU() ) self.s1 = MBConvStage(64, 64, num_layers=2) self.s2 = MBConvStage(64, 128, num_layers=6) self.s3 = TransformerStage(128, 256, num_layers=14) self.s4 = TransformerStage(256, 512, num_layers=2) def forward(self, x): x = self.s0(x) # /2 x = self.s1(x) # /4 x = self.s2(x) # /8 x = self.s3(x) # /16 x = self.s4(x) # /32 return x

2.2 关键组件实现

相对位置编码实现：

class RelPosEmb(nn.Module): def __init__(self, max_pos=14): super().__init__() self.rel_height = nn.Parameter(torch.randn(2 * max_pos - 1, 1)) self.rel_width = nn.Parameter(torch.randn(2 * max_pos - 1, 1)) def forward(self, q): B, H, W, _ = q.shape # 生成高度方向相对位置偏置 height_bias = self._get_rel_pos(self.rel_height, H) # 生成宽度方向相对位置偏置 width_bias = self._get_rel_pos(self.rel_width, W) return height_bias + width_bias

Transformer块集成：

class TransformerBlock(nn.Module): def __init__(self, dim, heads=8): super().__init__() self.norm1 = nn.LayerNorm(dim) self.attn = RelativeAttention(dim, heads) self.norm2 = nn.LayerNorm(dim) self.mlp = nn.Sequential( nn.Linear(dim, dim * 4), nn.SiLU(), nn.Linear(dim * 4, dim) ) def forward(self, x): x = x + self.attn(self.norm1(x)) x = x + self.mlp(self.norm2(x)) return x

3. 训练技巧与优化

3.1 数据增强策略

CoAtNet原文采用RandAugment增强组合：

• 随机裁剪（比例0.08-1.0）
• 水平翻转（概率0.5）
• 颜色抖动（亮度0.4，对比度0.4，饱和度0.4）
• MixUp（α=0.8）
• CutMix（α=1.0）

from torchvision import transforms train_transform = transforms.Compose([ transforms.RandomResizedCrop(224, scale=(0.08, 1.0)), transforms.RandomHorizontalFlip(), transforms.ColorJitter(0.4, 0.4, 0.4), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ])

3.2 学习率调度

采用余弦退火配合线性预热：

1. 前5个epoch线性增加学习率
2. 之后按余弦曲线衰减

from torch.optim.lr_scheduler import SequentialScheduler, LinearLR, CosineAnnealingLR optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3) scheduler = SequentialScheduler([ LinearLR(optimizer, start_factor=0.01, total_iters=5), CosineAnnealingLR(optimizer, T_max=95) ])

3.3 混合精度训练

使用AMP自动混合精度加速训练：

scaler = torch.cuda.amp.GradScaler() for inputs, targets in dataloader: 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. 性能对比与调优

4.1 不同配置对比实验

在CIFAR-100上的测试结果：

4.2 关键调优方向

1. 感受野平衡：

   • 浅层使用较大卷积核（5×5）
   • 深层减少注意力头数（4→2）

2. 通道维度调整：

   # 原配置 dims = [64, 128, 256, 512] # 优化配置（减少显存消耗） dims = [64, 96, 192, 384]

3. 注意力优化：

   • 局部窗口注意力（Swin风格）
   • 轴向注意力（分离行列计算）

class LocalAttention(nn.Module): def __init__(self, dim, window_size=7): super().__init__() self.window_size = window_size self.qkv = nn.Linear(dim, dim * 3) 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) # 分组 q, k, v = self.qkv(x).chunk(3, dim=-1) # 计算窗口内注意力 attn = (q @ k.transpose(-2,-1)) / math.sqrt(C) attn = attn.softmax(dim=-1) x = (attn @ v) # 还原空间结构 x = x.permute(0,1,3,2,4,5).reshape(B,H,W,C) return x

5. 部署优化实践

5.1 模型量化

model = coatnet_small(pretrained=True) quantized_model = torch.quantization.quantize_dynamic( model, {nn.Linear, nn.Conv2d}, dtype=torch.qint8 ) torch.save(quantized_model.state_dict(), 'quantized.pth')

5.2 ONNX导出

dummy_input = torch.randn(1, 3, 224, 224) torch.onnx.export( model, dummy_input, "coatnet.onnx", opset_version=13, input_names=["input"], output_names=["output"], dynamic_axes={ "input": {0: "batch"}, "output": {0: "batch"} } )

5.3 TensorRT加速

trtexec --onnx=coatnet.onnx \ --saveEngine=coatnet.engine \ --fp16 \ --workspace=4096

在NVIDIA V100上的推理性能：