yolov8 CBAM注意力机制

### 如何在YOLOv8中应用CBAM注意力机制 #### 实现概述 在YOLOv8中集成CBAM（Convolutional Block Attention Module）可以显著提升模型对于目标检测任务中的特征表达能力。通过引入通道和空间维度上的注意力建设，能够更精准地捕捉图像的关键区域并抑制不重要的部分。 #### 修改`__init__.py`与`conv.py` 为了使项目识别新的模块，在项目的初始化文件即`__init__.py`以及卷积层定义所在的`conv.py`里添加对CBAM的支持： ```python # __init__.py 文件内的改动如下所示： from .conv import * # 原有导入语句保持不变 __all__.append('CBAM') # 将'CBAM'追加至__all__ # conv.py 文件内则需新增CBAM类的实现代码片段： import torch.nn as nn class CBAM(nn.Module): def __init__(self, gate_channels, reduction_ratio=16, pool_types=['avg', 'max']): super(CBAM, self).__init__() self.ChannelGate = ChannelGate(gate_channels, reduction_ratio, pool_types) self.SpatialGate = SpatialGate() def forward(self, x): x_out = self.ChannelGate(x) x_out = self.SpatialGate(x_out) return x_out class ChannelGate(nn.Module): def __init__(self, gate_channels, reduction_ratio=16, pool_types=['avg', 'max']): super(ChannelGate, self).__init__() self.gate_channels = gate_channels self.mlp = nn.Sequential( Flatten(), nn.Linear(gate_channels, gate_channels // reduction_ratio), nn.ReLU(), nn.Linear(gate_channels // reduction_ratio, gate_channels) ) self.pool_types = pool_types self.avgpool = nn.AdaptiveAvgPool2d(1) self.maxpool = nn.AdaptiveMaxPool2d(1) def forward(self, x): channel_att_sum = None for pool_type in self.pool_types: if pool_type == 'avg': avg_pool = self.avgpool(x) channel_att_raw = self.mlp(avg_pool) elif pool_type == 'max': max_pool = self.maxpool(x) channel_att_raw = self.mlp(max_pool) if channel_att_sum is None: channel_att_sum = channel_att_raw else: channel_att_sum += channel_att_raw scale = torch.sigmoid(channel_att_sum).unsqueeze(2).unsqueeze(3).expand_as(x) return x * scale class SpatialGate(nn.Module): def __init__(self): super(SpatialGate, self).__init__() kernel_size = 7 self.compress = ChannelPool() self.spatial = BasicConv(2, 1, kernel_size, stride=1, padding=(kernel_size-1) // 2, relu=False) def forward(self, x): x_compress = self.compress(x) x_out = self.spatial(x_compress) scale = torch.sigmoid(x_out) # broadcasting return x * scale class ChannelPool(nn.Module): def forward(self, x): return torch.cat((torch.max(x, 1)[0].unsqueeze(1), torch.mean(x, 1).unsqueeze(1)), dim=1) class BasicConv(nn.Module): def __init__(self, in_planes, out_planes, kernel_size, stride=1, padding=0, dilation=1, groups=1, relu=True, bn=True, bias=False): super(BasicConv, self).__init__() self.out_channels = out_planes self.conv = nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias) self.bn = nn.BatchNorm2d(out_planes, eps=1e-5, momentum=0.01, affine=True) if bn else None self.relu = nn.ReLU() if relu else None def forward(self, x): x = self.conv(x) if self.bn is not None: x = self.bn(x) if self.relu is not None: x = self.relu(x) return x ``` 以上操作确保了框架能正确加载自定义组件，并允许后续配置文件调用此功能[^4]。 #### 调整`task.py`以支持CBAM 为了让训练过程利用上新加入的CBAM结构，需要编辑负责创建网络实例的任务脚本`task.py`: ```python def build_model(cfg, num_classes=None): model = Model(cfg, ch=3, nc=num_classes or cfg['nc']) # 遍历model.modules(), 对特定类型的layer插入CBAM for name, m in model.named_modules(): if isinstance(m, (nn.Conv2d)): # 可根据需求更改条件判断依据 setattr(model, name, nn.Sequential(m, CBAM(m.out_channels))) return model ``` 这段逻辑遍历整个神经网络图谱，针对满足一定条件的节点附加CBAM单元，从而完成对原有架构的功能扩展。 #### 更新配置文件`yolov8.yaml` 最后一步是在超参数设定文档`yolov8.yaml`里面指定哪些地方应该启用CBAM特性。这通常涉及到修改backbone或者neck部分的设计蓝图来反映这些变化： ```yaml # yolov8.yaml 的部分内容展示 depth_multiple: 0.33 # 模型深度系数 width_multiple: 0.50 # 模型宽度系数 head: # 头部设置... backbone: [[-1, 1, Focus], [-1, 1, Conv, [64]], ... ,[-1, 1, C3]] # 主干网路... neck: [[-1, 1, SPPF], ... ] # 特征金字塔... cbam_blocks: ['C3'] # 列表形式指明要添加CBAM的具体block名称 ``` 这里假设只给定名为'C3'的基础块加入了CBAM处理；实际部署时可根据个人喜好灵活调整。