1. 核心挑战
在深度学习实践中,我们经常遇到这样的问题:如何优化一个已经训练好的模型(如ResNet18)?
直接修改模型结构(如插入注意力模块)是否会破坏原有的特征提取能力?
如何制定训练策略,既能利用预训练权重,又能让新插入的模块快速学习?
2. 为什么可以将CBAM插入预训练模型?
通常认为,修改预训练模型的结构会导致权重失效。但在 ResNet 中插入 CBAM 模块是可行的,核心原因在于 CBAM 的初始化特性:
- 初始状态接近“直通”:
CBAM 模块的最终输出计算公式为:Output = Input * Sigmoid(Attention)。
在初始化阶段,卷积层和全连接层的权重接近于 0,导致 Attention 图的值也接近 0。
由于Sigmoid(0) = 0.5,因此在训练初期,CBAM 模块的操作近似于Input * 0.5。
- 保留特征结构:
Input * 0.5仅仅是对特征数值的线性缩放,完整保留了特征图的空间结构和相对关系。
这保证了下游的预训练层接收到的输入仍然是结构完好的特征,而不是混乱的噪声。
因此,我们可以将 CBAM “无缝注入”到预训练 ResNet 中,而不破坏其核心能力。
3. 模型架构:ResNet18 + CBAM
我们将 CBAM 模块插入到 ResNet18 的每一个layer(残差块组)之后。
3.1 代码实现
import torch import torch.nn as nn from torchvision import models class ResNet18_CBAM(nn.Module): def __init__(self, num_classes=10, pretrained=True, cbam_ratio=16, cbam_kernel=7): super().__init__() # 1. 加载预训练ResNet18 self.backbone = models.resnet18(pretrained=pretrained) # 2. 修改首层卷积以适应 CIFAR-10 的小尺寸输入 (32x32) # 原版 ResNet 针对 ImageNet (224x224),首层卷积核大且步长为2,会导致小图特征丢失 self.backbone.conv1 = nn.Conv2d( in_channels=3, out_channels=64, kernel_size=3, stride=1, padding=1, bias=False ) self.backbone.maxpool = nn.Identity() # 移除 MaxPool 层,保留更多空间信息 # 3. 在每个残差块组后添加 CBAM 模块 self.cbam_layer1 = CBAM(in_channels=64, ratio=cbam_ratio, kernel_size=cbam_kernel) self.cbam_layer2 = CBAM(in_channels=128, ratio=cbam_ratio, kernel_size=cbam_kernel) self.cbam_layer3 = CBAM(in_channels=256, ratio=cbam_ratio, kernel_size=cbam_kernel) self.cbam_layer4 = CBAM(in_channels=512, ratio=cbam_ratio, kernel_size=cbam_kernel) # 4. 修改分类头 self.backbone.fc = nn.Linear(in_features=512, out_features=num_classes) def forward(self, x): # Stem 层 x = self.backbone.conv1(x) x = self.backbone.bn1(x) x = self.backbone.relu(x) # Layer 1 + CBAM x = self.backbone.layer1(x) x = self.cbam_layer1(x) # Layer 2 + CBAM x = self.backbone.layer2(x) x = self.cbam_layer2(x) # Layer 3 + CBAM x = self.backbone.layer3(x) x = self.cbam_layer3(x) # Layer 4 + CBAM x = self.backbone.layer4(x) x = self.cbam_layer4(x) # 分类头 x = self.backbone.avgpool(x) x = torch.flatten(x, 1) x = self.backbone.fc(x) return x4. 训练策略:三阶段渐进式微调
为了平衡“保留预训练知识”和“学习新模块”的需求,我们采用了差异化学习率和分阶段解冻的策略。可以将模型看作一个公司团队:
- 预训练层(ResNet):资深专家,经验丰富,只需微调。
- 新模块(CBAM/FC):新入职实习生,一张白纸,需要快速学习。
4.1 阶段一:预热实习生 (Epoch 1-5)
- 解冻对象:仅
CBAM模块和分类头 (fc)。 - 冻结对象:所有 ResNet 主干层。
- 学习率:
1e-3(高学习率)。 - 目标:让新模块快速学习如何配合预训练特征工作,建立初步的分类边界和注意力机制。
4.2 阶段二:唤醒高层专家 (Epoch 6-20)
- 解冻对象:增加解冻
layer3和layer4(高层卷积)。 - 冻结对象:
conv1,layer1,layer2(底层卷积)保持冻结。 - 学习率:
1e-4(中等学习率)。 - 原理:高层网络学习“构图和概念”(如车轮、猫脸),与具体任务强相关,需要重新适应 CIFAR-10。而底层网络学习“笔触和纹理”(如边缘、颜色),是通用的,暂时不需要变动。
4.3 阶段三:全员协同微调 (Epoch 21-50)
- 解冻对象:所有层。
- 学习率:
1e-5(低学习率)。 - 目标:进行端到端的全局微调,让底层特征也做微小的调整以完美适配新任务。
5. 核心训练代码
实现上述策略的关键函数:
def set_trainable_layers(model, trainable_parts): print(f"\n---> 解冻以下部分并设为可训练: {trainable_parts}") for name, param in model.named_parameters(): param.requires_grad = False # 先默认全冻结 for part in trainable_parts: if part in name: # 如果参数名包含指定部分,则解冻 param.requires_grad = True break def train_staged_finetuning(model, criterion, train_loader, test_loader, device, epochs): # ... (省略初始化代码) for epoch in range(1, epochs + 1): # --- 动态调整策略 --- if epoch == 1: print("阶段 1:训练注意力模块和分类头") set_trainable_layers(model, ["cbam", "backbone.fc"]) # 只优化 requires_grad=True 的参数 optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=1e-3) elif epoch == 6: print("阶段 2:解冻高层卷积层") set_trainable_layers(model, ["cbam", "backbone.fc", "backbone.layer3", "backbone.layer4"]) optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=1e-4) elif epoch == 21: print("阶段 3:全局微调") for param in model.parameters(): param.requires_grad = True optimizer = optim.Adam(model.parameters(), lr=1e-5) # ... (后续标准训练与测试循环)6. 实验结果
通过这种精细的训练策略,模型在 CIFAR-10 数据集上最终达到了90.15%的测试准确率。这证明了:
- CBAM 能够有效增强特征提取能力。
- 分阶段微调策略能够有效防止预训练权重的破坏,同时让新模块充分收敛。