当前位置:   article > 正文

b站B导的yoloV7版本添加注意力机制_yolov7添加注意力机制

yolov7添加注意力机制

视频参考:

【Pytorch 图像处理中注意力机制的代码详解与应用(Bubbliiiing 深度学习 教程)】

yolov7博客参考:

睿智的目标检测61——Pytorch搭建YoloV7目标检测平台

注意力机制介绍:

神经网络学习小记录64——Pytorch 图像处理中注意力机制的解析与代码详解

步骤:

①获取注意力机制的代码

https://github.com/bubbliiiing/yolov4-tiny-pytorch/blob/master/nets/attention.py

②看b导的yolov7博客介绍以及视频

看完视频可以知道,添加注意力机制最好选在加强特征提取网络。注意力机制是一个即插即用的模块,理论上可以放在任何一个特征层后面,可以放在主干网络,也可以放在加强特征提取网络。由于放置在主干会导致网络的预训练权重无法使用,所以将注意力机制应用加强特征提取网络上。

③yolov7添加注意力的地方:

一共五个地方,如图所示,三个feat和两个上采样的地方。

④添加注意力机制的代码,放在net文件夹下

⑤在nets/yolo.py增加代码

添加模块

  1. from .attention import se_block,cbam_block,eca_block
  2. attention_blocks = [se_block,cbam_block,eca_block]

找到yolo_body(快捷键ctrl+f搜索yolo_body)

在参数列表添加attention参数,attention = 0:表示原版,attention=1表示senet,attention=2表示cbam,attention=3表示ecanet,这个序号与我们在文件头部定义的

attention_blocks = [se_block,cbam_block,eca_block]有关。

在该__init__函数的末尾加上:

  1. self.attention = attention
  2. if attention >= 1 and attention <= 3:
  3. self.feat1_attention = attention_blocks[attention - 1](512)#括号表示通道数
  4. self.feat2_attention = attention_blocks[attention - 1](1024)
  5. self.feat3_attention = attention_blocks[attention - 1](1024)
  6. self.upsample_attention5 = attention_blocks[attention - 1](256)
  7. self.upsample_attention4 = attention_blocks[attention - 1](128)

在forward函数增加如下代码:

  1. if self.attention >= 1 and self.attention <= 3:
  2. feat1 = self.feat1_attention(feat1)
  3. feat2 = self.feat2_attention(feat2)
  4. feat3 = self.feat3_attention(feat3)
  1. if self.attention >= 1 and self.attention <= 3:
  2. P5_upsample = self.upsample_attention5(P5_upsample)
  3. if self.attention >= 1 and self.attention <= 3:
  4. P4_upsample = self.upsample_attention4(P4_upsample)

接下来只需要修改nets/yolo.py中yolo_body初始化参数列表的attention的值即可切换不同的注意力机制

完成。

本文内容由网友自发贡献,转载请注明出处:https://www.wpsshop.cn/w/IT小白/article/detail/346258
推荐阅读
  

闽ICP备14008679号