yolov5项目增加注意力机制

风吹斜刘海2025-03-18 18:58

1. 添加步骤

1.1 models/common.py 加入新增的网络结构

1.2 在models/yolo.py注册组件

1.3 models/yolov5_xxx.yaml 新建一个文件夹，修改模型结构配置

1.3.1 注意力模块的添加位置

注意力机制可以添加在backbone，Neck，Head等部分，常见的有两种，一种是在主干的SPPF前添加一层；二是将backbone中的C3全部替换。不同的位置效果可能不同，需要我们用数据集反复测试。

这就是所谓的调参

1.3.2 修改后续层的参数

当在网络中添加了新的层之后，该层后续层的编号会发生变化。

1. Concat

如原本Concat层的系数也要修改，[-1, 14],[-1, 10] -> [-1, 15], [-1, 11]

2. Detect

同样Detec指定的是[17, 20, 23]层，在添加了SE注意力层之后，就需要对Detect进行修改，[17, 20, 23] -> [18, 21, 24]

1.4 验证是否加入成功

1.5 train.py 修改'--cfg'默认参数，训练时指定模型结构配置文件

2. 增加SE注意力机制

2.1 新增SE模块

放在common.py中

2.2 注册组件

yolo.py

2.3 创建yolov5s_SE.yaml文件

2.3.1 将SE模块写入第9层

2.3.2 修改后续层的参数

Concat

改为

Detect

更改为

3. 验证是否加入成功

修改yolo.py

运行yolo.py 第9层成功加入SE模块