YOLOv5s添加CBAM注意力机制

### 如何在 YOLOv5s 模型中添加 CBAM 注意力机制 要在 YOLOv5s 中集成 CBAM (Convolutional Block Attention Module)，可以通过修改模型配置文件以及实现 CBAM 的代码来完成。以下是具体方法： #### 修改模型配置文件 YOLOv5 使用 YAML 文件定义模型结构。为了将 CBAM 集成到 YOLOv5s，可以在 `models/yolov5s.yaml` 或自定义的 `.yaml` 文件中调整网络架构。 假设我们希望在 backbone 和 head 层都加入 CBAM，则可以按照以下方式修改配置文件[^1]： ```yaml # yolov5s_cbam.yaml nc: 80 # 类别数 depth_multiple: 0.33 # 深度倍率 width_multiple: 0.25 # 宽度倍率 # Backbone with CBAM backbone: # [from, number, module, args] [[-1, 1, Focus, [64, 3]], # 0-PADDED/STRIDED CONVOLUTION [-1, 1, Conv, [64, 3, 1]], # 1-BASIC CONVOLUTION LAYER [-1, 1, C3, [64, 1]], # 2-CSP BOTTLENECK WITH 3 CONVS [-1, 1, CBAM, []], # 添加CBAM模块 [-1, 3, BottleneckCSP, [128, True]], # CSPBLOCKS WITH DOWN SAMPLING ... ] # Head with CBAM head: [[-1, 1, SPP, [768, [5, 9, 13]]], [-1, 1, CBAM, []], # 在Head部分也添加CBAM [-1, 3, BottleneckCSP, [768]], ... ] ``` 以上示例展示了如何通过简单的结构调整在网络的不同阶段引入 CBAM 模块。 --- #### 实现 CBAM 模块 接下来需要编写 CBAM 的 PyTorch 实现并将其嵌入到 YOLOv5 的代码库中。以下是完整的 CBAM 模块代码： ```python import torch import torch.nn as nn 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.fc1 = nn.Conv2d(in_planes, in_planes // ratio, 1, bias=False) self.relu1 = nn.ReLU() self.fc2 = nn.Conv2d(in_planes // ratio, in_planes, 1, bias=False) self.sigmoid = nn.Sigmoid() def forward(self, x): avg_out = self.fc2(self.relu1(self.fc1(self.avg_pool(x)))) max_out = self.fc2(self.relu1(self.fc1(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__() assert kernel_size in (3, 7), 'kernel size must be 3 or 7' padding = 3 if kernel_size == 7 else 1 self.conv1 = nn.Conv2d(2, 1, kernel_size, padding=padding, 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, planes, ratio=16, kernel_size=7): super(CBAM, self).__init__() self.channel_attention = ChannelAttention(planes, ratio) self.spatial_attention = SpatialAttention(kernel_size) def forward(self, x): x = x * self.channel_attention(x) # Apply channel attention first. x = x * self.spatial_attention(x) # Then apply spatial attention. return x ``` 此代码实现了 CBAM 的两个子模块——通道注意力和空间注意力，并组合为最终的 CBAM 模块。 --- #### 将 CBAM 嵌入到 YOLOv5 要让 YOLOv5 支持 CBAM，需将上述 CBAM 模块保存至 `models/common.py` 并注册到 YOLOv5 的构建函数中。例如，在 `common.py` 中新增如下内容： ```python def CBAM(c1, c2=None, *, ratio=16, k=7): """CBAM implementation.""" c2 = c2 if c2 is not None else c1 return CBAMModule(c2, ratio=ratio, kernel_size=k) ``` 随后更新 `detect.py`, `train.py` 等脚本中的导入路径以支持新模块。 --- #### 训练与验证 最后一步是在训练过程中加载新的配置文件 `yolov5s_cbam.yaml` 运行实验。命令如下所示： ```bash python train.py --img 640 --batch 16 --epochs 50 --data coco128.yaml --cfg models/yolov5s_cbam.yaml --weights yolov5s.pt ``` 这会基于预训练权重启动带有 CBAM 的 YOLOv5s 模型训练过程。 ---