PyTorch实战：手把手教你给ResNet加上SENet、SKNet和CBAM注意力模块（附完整代码）

PyTorch实战：手把手教你给ResNet加上SENet、SKNet和CBAM注意力模块（附完整代码）

在计算机视觉领域，注意力机制已经成为提升模型性能的重要工具。本文将带你深入实践，一步步为ResNet模型集成三种主流注意力模块：SENet、SKNet和CBAM。无论你是想快速验证这些模块的效果，还是需要在项目中灵活应用，这里提供的完整代码和实战技巧都能让你事半功倍。

1. 准备工作与环境配置

在开始之前，我们需要确保开发环境准备就绪。推荐使用Python 3.8+和PyTorch 1.8+版本，这些版本对注意力机制的支持最为完善。

首先安装必要的依赖库：

pip install torch torchvision matplotlib tqdm

对于硬件配置，虽然这些注意力模块会增加少量计算量，但现代GPU都能很好地支持。以下是不同规模模型的大致显存需求：

提示：在添加注意力模块后，显存占用通常会增加10%-20%，训练速度会降低15%-30%，但模型精度往往能有显著提升。

2. 基础ResNet模型解析

在添加注意力模块前，我们需要理解ResNet的基本结构。ResNet的核心是残差块（Residual Block），它通过跳跃连接解决了深层网络训练困难的问题。

一个典型的BasicBlock结构如下：

class BasicBlock(nn.Module): expansion = 1 def __init__(self, inplanes, planes, stride=1, downsample=None): super(BasicBlock, self).__init__() self.conv1 = conv3x3(inplanes, planes, stride) self.bn1 = nn.BatchNorm2d(planes) self.relu = nn.ReLU(inplace=True) self.conv2 = conv3x3(planes, planes) self.bn2 = nn.BatchNorm2d(planes) self.downsample = downsample self.stride = stride def forward(self, x): residual = x out = self.conv1(x) out = self.bn1(out) out = self.relu(out) out = self.conv2(out) out = self.bn2(out) if self.downsample is not None: residual = self.downsample(x) out += residual out = self.relu(out) return out

我们将在这个基础结构上添加不同的注意力模块。添加位置的选择很重要：

• 浅层添加：更适合捕捉纹理等低级特征
• 深层添加：更适合处理语义等高级特征
• 每个残差块后添加：全面增强特征表达能力

3. 集成SENet注意力模块

SENet(Squeeze-and-Excitation Network)是最早的通道注意力机制之一，它通过学习通道间的重要性来增强特征表示。

3.1 SENet模块实现

class SELayer(nn.Module): def __init__(self, channel, reduction=16): super(SELayer, self).__init__() self.avg_pool = nn.AdaptiveAvgPool2d(1) self.fc = nn.Sequential( nn.Linear(channel, channel // reduction, bias=False), nn.ReLU(inplace=True), nn.Linear(channel // reduction, channel, bias=False), nn.Sigmoid() ) def forward(self, x): b, c, _, _ = x.size() y = self.avg_pool(x).view(b, c) y = self.fc(y).view(b, c, 1, 1) return x * y.expand_as(x)

3.2 集成到ResNet

将SENet集成到ResNet的BasicBlock中：

class SEBasicBlock(nn.Module): expansion = 1 def __init__(self, inplanes, planes, stride=1, downsample=None, reduction=16): super(SEBasicBlock, self).__init__() self.conv1 = conv3x3(inplanes, planes, stride) self.bn1 = nn.BatchNorm2d(planes) self.relu = nn.ReLU(inplace=True) self.conv2 = conv3x3(planes, planes) self.bn2 = nn.BatchNorm2d(planes) self.se = SELayer(planes, reduction) self.downsample = downsample self.stride = stride def forward(self, x): residual = x out = self.conv1(x) out = self.bn1(out) out = self.relu(out) out = self.conv2(out) out = self.bn2(out) out = self.se(out) # 添加SE模块 if self.downsample is not None: residual = self.downsample(x) out += residual out = self.relu(out) return out

注意：reduction参数控制压缩比例，通常设置为16，但在小模型上可以尝试8或4以减少信息损失。

4. 集成SKNet注意力模块

SKNet(Selective Kernel Network)通过动态选择不同大小的卷积核来适应不同尺度的特征。

4.1 SKNet模块实现

class SKConv(nn.Module): def __init__(self, features, M=2, G=32, r=16, stride=1, L=32): super(SKConv, self).__init__() d = max(int(features/r), L) self.M = M self.features = features self.convs = nn.ModuleList([]) for i in range(M): self.convs.append(nn.Sequential( nn.Conv2d(features, features, kernel_size=3+i*2, stride=stride, padding=1+i, groups=G), nn.BatchNorm2d(features), nn.ReLU(inplace=False) )) self.fc = nn.Linear(features, d) self.fcs = nn.ModuleList([]) for i in range(M): self.fcs.append(nn.Linear(d, features)) self.softmax = nn.Softmax(dim=1) def forward(self, x): for i, conv in enumerate(self.convs): fea = conv(x).unsqueeze_(dim=1) if i == 0: feas = fea else: feas = torch.cat([feas, fea], dim=1) fea_U = torch.sum(feas, dim=1) fea_s = fea_U.mean(-1).mean(-1) fea_z = self.fc(fea_s) for i, fc in enumerate(self.fcs): vector = fc(fea_z).unsqueeze_(dim=1) if i == 0: attention_vectors = vector else: attention_vectors = torch.cat([attention_vectors, vector], dim=1) attention_vectors = self.softmax(attention_vectors) attention_vectors = attention_vectors.unsqueeze(-1).unsqueeze(-1) fea_v = (feas * attention_vectors).sum(dim=1) return fea_v

4.2 集成到ResNet

class SKBasicBlock(nn.Module): expansion = 1 def __init__(self, inplanes, planes, stride=1, downsample=None, M=2, G=32, r=16): super(SKBasicBlock, self).__init__() self.conv1 = conv3x3(inplanes, planes, stride) self.bn1 = nn.BatchNorm2d(planes) self.relu = nn.ReLU(inplace=True) self.sk = SKConv(planes, M, G, r) self.conv2 = conv3x3(planes, planes) self.bn2 = nn.BatchNorm2d(planes) self.downsample = downsample self.stride = stride def forward(self, x): residual = x out = self.conv1(x) out = self.bn1(out) out = self.relu(out) out = self.sk(out) # 添加SK模块 out = self.conv2(out) out = self.bn2(out) if self.downsample is not None: residual = self.downsample(x) out += residual out = self.relu(out) return out

5. 集成CBAM注意力模块

CBAM(Convolutional Block Attention Module)结合了通道注意力和空间注意力，是一种更全面的注意力机制。

5.1 CBAM模块实现

class ChannelAttention(nn.Module): def __init__(self, in_planes, ratio=16): super(ChannelAttention, self).__init__() self.avg_pool = nn.AdaptiveAvgPool2d(1) self.max_pool = nn.AdaptiveMaxPool2d(1) self.fc = nn.Sequential( nn.Conv2d(in_planes, in_planes // ratio, 1, bias=False), nn.ReLU(), nn.Conv2d(in_planes // ratio, in_planes, 1, bias=False) ) self.sigmoid = nn.Sigmoid() def forward(self, x): avg_out = self.fc(self.avg_pool(x)) max_out = self.fc(self.max_pool(x)) out = avg_out + max_out return self.sigmoid(out) class SpatialAttention(nn.Module): def __init__(self, kernel_size=7): super(SpatialAttention, self).__init__() self.conv1 = nn.Conv2d(2, 1, kernel_size, padding=kernel_size//2, bias=False) self.sigmoid = nn.Sigmoid() def forward(self, x): avg_out = torch.mean(x, dim=1, keepdim=True) max_out, _ = torch.max(x, dim=1, keepdim=True) x = torch.cat([avg_out, max_out], dim=1) x = self.conv1(x) return self.sigmoid(x) class CBAM(nn.Module): def __init__(self, in_planes, ratio=16, kernel_size=7): super(CBAM, self).__init__() self.ca = ChannelAttention(in_planes, ratio) self.sa = SpatialAttention(kernel_size) def forward(self, x): x = self.ca(x) * x x = self.sa(x) * x return x

5.2 集成到ResNet

class CBAMBasicBlock(nn.Module): expansion = 1 def __init__(self, inplanes, planes, stride=1, downsample=None, ratio=16, kernel_size=7): super(CBAMBasicBlock, self).__init__() self.conv1 = conv3x3(inplanes, planes, stride) self.bn1 = nn.BatchNorm2d(planes) self.relu = nn.ReLU(inplace=True) self.conv2 = conv3x3(planes, planes) self.bn2 = nn.BatchNorm2d(planes) self.cbam = CBAM(planes, ratio, kernel_size) self.downsample = downsample self.stride = stride def forward(self, x): residual = x out = self.conv1(x) out = self.bn1(out) out = self.relu(out) out = self.conv2(out) out = self.bn2(out) out = self.cbam(out) # 添加CBAM模块 if self.downsample is not None: residual = self.downsample(x) out += residual out = self.relu(out) return out

6. 实验对比与调优建议

在实际应用中，不同的注意力模块适用于不同的场景。以下是一些调优建议：

1. 模块选择：

   • 对于计算资源有限的场景，优先考虑SENet
   • 需要处理多尺度特征时，SKNet表现更佳
   • 追求最高精度时，CBAM通常是更好的选择

2. 超参数设置：

   • reduction ratio：通常16是好的起点，小模型可以尝试8
   • SKNet的M值：2-3个分支足够，更多分支收益递减
   • CBAM的空间注意力核大小：7x7在大多数情况下效果最好

3. 训练技巧：

   • 初始学习率可以比标准ResNet小10%-20%
   • 使用warmup策略有助于稳定训练
   • 注意力模块的参数可以使用稍大的权重衰减(1e-4)

以下是在CIFAR-100上的对比实验结果：

7. 完整代码示例

以下是集成CBAM的完整ResNet实现：

import torch import torch.nn as nn import torch.nn.functional as F def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1): return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=dilation, groups=groups, bias=False, dilation=dilation) class CBAMResNet(nn.Module): def __init__(self, block, layers, num_classes=1000, zero_init_residual=False, groups=1, width_per_group=64, replace_stride_with_dilation=None, norm_layer=None): super(CBAMResNet, self).__init__() if norm_layer is None: norm_layer = nn.BatchNorm2d self._norm_layer = norm_layer self.inplanes = 64 self.dilation = 1 if replace_stride_with_dilation is None: replace_stride_with_dilation = [False, False, False] self.groups = groups self.base_width = width_per_group self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=7, stride=2, padding=3, bias=False) self.bn1 = norm_layer(self.inplanes) self.relu = nn.ReLU(inplace=True) self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) self.layer1 = self._make_layer(block, 64, layers[0]) self.layer2 = self._make_layer(block, 128, layers[1], stride=2, dilate=replace_stride_with_dilation[0]) self.layer3 = self._make_layer(block, 256, layers[2], stride=2, dilate=replace_stride_with_dilation[1]) self.layer4 = self._make_layer(block, 512, layers[3], stride=2, dilate=replace_stride_with_dilation[2]) self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) self.fc = nn.Linear(512 * block.expansion, num_classes) for m in self.modules(): if isinstance(m, nn.Conv2d): nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)): nn.init.constant_(m.weight, 1) nn.init.constant_(m.bias, 0) if zero_init_residual: for m in self.modules(): if isinstance(m, CBAMBasicBlock): nn.init.constant_(m.bn2.weight, 0) def _make_layer(self, block, planes, blocks, stride=1, dilate=False): norm_layer = self._norm_layer downsample = None previous_dilation = self.dilation if dilate: self.dilation *= stride stride = 1 if stride != 1 or self.inplanes != planes * block.expansion: downsample = nn.Sequential( conv1x1(self.inplanes, planes * block.expansion, stride), norm_layer(planes * block.expansion), ) layers = [] layers.append(block(self.inplanes, planes, stride, downsample, self.groups, self.base_width, previous_dilation, norm_layer)) self.inplanes = planes * block.expansion for _ in range(1, blocks): layers.append(block(self.inplanes, planes, groups=self.groups, base_width=self.base_width, dilation=self.dilation, norm_layer=norm_layer)) return nn.Sequential(*layers) def forward(self, x): x = self.conv1(x) x = self.bn1(x) x = self.relu(x) x = self.maxpool(x) x = self.layer1(x) x = self.layer2(x) x = self.layer3(x) x = self.layer4(x) x = self.avgpool(x) x = torch.flatten(x, 1) x = self.fc(x) return x

在实际项目中，我发现CBAM模块在图像分类任务上表现最为稳定，而SKNet在目标检测任务中因其多尺度特性往往能有更好的表现。对于资源受限的部署环境，经过适当剪枝的SENet模型是更经济的选择。