主题:在 YOLOv11 目标检测模型中集成 CBAM 注意力机制,并在 Kaggle 上训练自定义数据集
一、学习目标
-
理解 CBAM(卷积块注意力模块)的原理与作用。
-
将 CBAM 模块成功嵌入 YOLOv11 的网络结构中。
-
在 Kaggle 平台(云端 GPU)上完成模型训练,并分析效果。
二、上午:环境准备与代码整合
2.1 遇到的困难
-
本地代码(Windows)可以正常运行,但迁移到 Kaggle Notebook 后反复出现
KeyError: 'CBAM'和ImportError。 -
原因:Kaggle 上的 Ultralytics 库是原装的,未像本地那样修改过
tasks.py文件,导致parse_model函数不认识 CBAM 模块。
2.2 解决方案
-
不修改库源码,而是在 Notebook 中动态定义 CBAM 类,并手动修改 YAML 配置文件,为 CBAM 层显式指定输入通道数。
-
关键步骤:
-
定义
ChannelAttention、SpatialAttention和CBAM类。 -
将
CBAM注册到ultralytics.nn.tasks模块中:tasks.CBAM = CBAM。 -
读取原始 YAML 文件,找到 head 部分的
CBAM层,将其参数[]改为[256](根据网络结构确定通道数)。 -
同时修改 YAML 中的
nc为数据集的真实类别数(6类)。 -
生成修正后的 YAML 文件供训练使用。
-
2.3 数据集配置
-
数据集路径:
/kaggle/input/.../coal -
类别(6类):
miner,helmet,vest,mask,fire,face -
编写脚本自动检测
train和val图片目录,生成数据集 YAML。
三、下午:训练模型与结果分析
3.1 训练配置
-
模型:YOLOv11n + CBAM(nano 版本)
-
轮数:100 epochs
-
图片尺寸:640×640
-
批次大小:16
-
优化器:默认 SGD,学习率 0.01
3.2 训练结果
text
Class Images Instances Box(P R mAP50 mAP50-95
all 531 1555 0.715 0.615 0.644 0.401
miner 186 319 0.666 0.537 0.592 0.360
helmet 190 292 0.685 0.468 0.512 0.242
vest 155 221 0.637 0.498 0.521 0.288
mask 104 133 0.814 0.729 0.750 0.539
fire 91 229 0.786 0.817 0.830 0.596
face 163 361 0.705 0.641 0.657 0.383
-
整体 mAP50-95 = 0.401,效果不够理想。
-
表现较好的类:mask(0.539)、fire(0.596)
-
表现较差的类:helmet(0.242)、vest(0.288)
3.3 原因分析
-
数据集较小:训练图片可能只有几百张,实例数不均衡(mask 仅133个,face 有361个)。
-
模型容量不足:YOLOv11n 参数量仅 259 万,特征提取能力有限,尤其对 helmet、vest 这类小目标。
-
训练超参数可能不是最优:100 轮可能不够收敛,默认数据增强对小目标帮助有限。
-
CBAM 提升有限:注意力机制通常只能带来 1~3% 的提升,无法弥补数据量不足的根本问题。
四、改进思路(未执行,作为笔记)
| 方法 | 预期提升 | 操作难度 |
|---|---|---|
| 增加训练轮数至 300 | 5~10% | 低 |
| 降低学习率 + 余弦退火 | 5~15% | 低 |
| 调整数据增强(mixup, copy_paste) | 3~8% | 低 |
| 换用 YOLOv11s 模型 | 15~25% | 中 |
| 收集更多训练图片(尤其是小目标) | 20~50% | 高 |
五、今日总结与反思
5.1 技术收获
-
学会了如何在不修改库源码的情况下,为 Ultralytics 框架添加自定义模块(CBAM)。
-
掌握了 YAML 配置文件的动态修改技巧(用 Python 读取、修改、写入)。
-
理解了模型参数量、数据集规模与最终精度之间的关系。
5.2 不足与教训
-
前期花费大量时间在"修补
parse_model"上,走了弯路。最终发现手动修改 YAML 参数更简单可靠。 -
训练前没有先做一个快速的小规模实验(如 20 轮),导致 100 轮跑完后才发现效果不佳。下次应先调参再全量训练。
-
对数据集本身的认知不足:没有事先分析类别分布、目标尺寸,导致对 helmet 等小目标预期过高。
5.3 下一步计划
-
尝试用改进后的超参数重新训练 300 轮。
-
如果效果仍不佳,换用 YOLOv11s 版本。
-
考虑使用数据增强策略( mosaic + mixup + copy_paste )强化小目标学习。