1. 项目概述
今天我想分享一个在YOLOv11目标检测模型中集成CBAM注意力机制的完整实践过程。作为一名计算机视觉工程师,我最近在Kaggle平台上完成了一个自定义数据集训练项目,目标是检测煤矿安全场景中的6类目标(矿工、安全帽、反光背心、口罩、火焰和人脸)。这个项目让我深刻体会到模型调优的复杂性,也积累了一些宝贵的实战经验。
CBAM(Convolutional Block Attention Module)是一种轻量级的注意力机制模块,包含通道注意力和空间注意力两部分。它能帮助模型更有效地聚焦于图像中的重要区域,对于煤矿安全这种复杂场景的目标检测尤为有用。不过在实际集成过程中,我遇到了不少意料之外的挑战,特别是在Kaggle环境下的代码适配问题。
需要模型API调用? 免费领10W Token,多模型网关一键接入 Claude、DeepSeek 等主流模型。
2. 环境准备与代码整合
2.1 CBAM模块原理与实现
CBAM的核心思想是通过两个连续的注意力机制(通道注意力和空间注意力)来增强特征表示。通道注意力通过全局平均池化和最大池化来学习通道间的重要性,而空间注意力则通过沿通道轴的平均和最大操作来学习空间位置的重要性。
在实现上,我采用了以下类结构:
python复制class ChannelAttention(nn.Module):
def __init__(self, in_planes, ratio=16):
super().__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))
return self.sigmoid(avg_out + max_out) * x
class SpatialAttention(nn.Module):
def __init__(self, kernel_size=7):
super().__init__()
self.conv = nn.Conv2d(2, 1, kernel_size, padding=kernel_size//2, bias=False)
self.sigmoid = nn.Sigmoid()
def forward(self, x):
a
