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 注册组件
2.3 创建yolov5s_SE.yaml文件
2.3.1 将SE模块写入第9层
2.3.2 修改后续层的参数
Concat
改为
Detect
更改为
3. 验证是否加入成功
运行yolo.py 第9层成功加入SE模块