### 如何在 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 模型训练过程。
---
最低0.47元/天 解锁文章
扫一扫
2801
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
