文章详情

短信预约-IT技能 免费直播动态提醒

请输入下面的图形验证码

提交验证

短信预约提醒成功

改进YOLOv5:添加EMA注意力机制

2023-09-05 16:14

关注

提示:文章写完后,目录可以自动生成,如何生成可参考右边的帮助文档

文章目录


本文主要介绍一种在YOLOv5-7.0中添加EMA注意力机制的方法。EMA注意力机制原论文地址,有关EMA注意力机制的解读可参考文章

在yolov5的models文件中新建一个名为EMA.py文件,将下述代码复制到EMA.py文件中并保存。

import torchfrom torch import nnclass EMA(nn.Module):    def __init__(self, channels, factor=8):        super(EMA, self).__init__()        self.groups = factor        assert channels // self.groups > 0        self.softmax = nn.Softmax(-1)        self.agp = nn.AdaptiveAvgPool2d((1, 1))        self.pool_h = nn.AdaptiveAvgPool2d((None, 1))        self.pool_w = nn.AdaptiveAvgPool2d((1, None))        self.gn = nn.GroupNorm(channels // self.groups, channels // self.groups)        self.conv1x1 = nn.Conv2d(channels // self.groups, channels // self.groups, kernel_size=1, stride=1, padding=0)        self.conv3x3 = nn.Conv2d(channels // self.groups, channels // self.groups, kernel_size=3, stride=1, padding=1)    def forward(self, x):        b, c, h, w = x.size()        group_x = x.reshape(b * self.groups, -1, h, w)  # b*g,c//g,h,w        x_h = self.pool_h(group_x)        x_w = self.pool_w(group_x).permute(0, 1, 3, 2)        hw = self.conv1x1(torch.cat([x_h, x_w], dim=2))        x_h, x_w = torch.split(hw, [h, w], dim=2)        x1 = self.gn(group_x * x_h.sigmoid() * x_w.permute(0, 1, 3, 2).sigmoid())        x2 = self.conv3x3(group_x)        x11 = self.softmax(self.agp(x1).reshape(b * self.groups, -1, 1).permute(0, 2, 1))        x12 = x2.reshape(b * self.groups, c // self.groups, -1)  # b*g, c//g, hw        x21 = self.softmax(self.agp(x2).reshape(b * self.groups, -1, 1).permute(0, 2, 1))        x22 = x1.reshape(b * self.groups, c // self.groups, -1)  # b*g, c//g, hw        weights = (torch.matmul(x11, x12) + torch.matmul(x21, x22)).reshape(b * self.groups, 1, h, w)        return (group_x * weights.sigmoid()).reshape(b, c, h, w)

1.导入EMA.py

在yolo.py文件开头导入EMA.py,代码如下:

from models.EMA import EMA

代码放在yolo.py位置如下图所示:
在这里插入图片描述

2.修改parse_model

这里主要是添加通道参数,再添加一个elif,把EMA添加进去,代码如下:

 elif m is EMA:               args = [ch[f], *args]

添加上述代码的位置可参考下图:
在这里插入图片描述


# YOLOv5 🚀 by Ultralytics, GPL-3.0 license# Parametersnc: 80  # number of classesdepth_multiple: 0.33  # model depth multiplewidth_multiple: 0.50  # layer channel multipleanchors:  - [10,13, 16,30, 33,23]  # P3/8  - [30,61, 62,45, 59,119]  # P4/16  - [116,90, 156,198, 373,326]  # P5/32# YOLOv5 v6.0 backbonebackbone:  # [from, number, module, args]  [[-1, 1, Conv, [64, 6, 2, 2]],  # 0-P1/2   [-1, 1, Conv, [128, 3, 2]],  # 1-P2/4   [-1, 3, C3, [128]],   [-1, 1, Conv, [256, 3, 2]],  # 3-P3/8   [-1, 6, C3, [256]],   [-1, 1, Conv, [512, 3, 2]],  # 5-P4/16   [-1, 9, C3, [512]],   [-1, 1, Conv, [1024, 3, 2]],  # 7-P5/32   [-1, 3, C3, [1024]],   [-1, 1, EMA, [8]],   [-1, 1, SPPF, [1024, 5]],  # 9  ]# YOLOv5 v6.0 headhead:  [[-1, 1, Conv, [512, 1, 1]],   [-1, 1, nn.Upsample, [None, 2, 'nearest']],   [[-1, 6], 1, Concat, [1]],  # cat backbone P4   [-1, 3, C3, [512, False]],  # 13   [-1, 1, Conv, [256, 1, 1]],   [-1, 1, nn.Upsample, [None, 2, 'nearest']],   [[-1, 4], 1, Concat, [1]],  # cat backbone P3   [-1, 3, C3, [256, False]],  # 17 (P3/8-small)   [-1, 1, Conv, [256, 3, 2]],   [[-1, 15], 1, Concat, [1]],  # cat head P4   [-1, 3, C3, [512, False]],  # 20 (P4/16-medium)   [-1, 1, Conv, [512, 3, 2]],   [[-1, 11], 1, Concat, [1]],  # cat head P5   [-1, 3, C3, [1024, False]],  # 23 (P5/32-large)   [[18, 21, 24], 1, Detect, [nc, anchors]],  # Detect(P3, P4, P5)  ]

上述代码将EMA注意力机制模块加在backbone层中最后C3模块后面,SPPF模块前面,仅供参考,具体添加位置要根据个人数据集的不同合理的添加。

[-1, 1, EMA, [8]], #-1代表连接上一层通道数,1是个数,8是EMA所需的参数(factor=8)

说明:因为在yolo.py文件parse_model函数中修改了通道参数,因此在yaml文件中无需添加通道参数,只需添加EMA函数所需的其他参数。在backbone中添加一层注意力机制模块,因此后续的层数都要加一,在head层中做如下改动。

[[-1, 15], 1, Concat, [1]],  #未改动前的第14层,在经过上述改动后改为15[[-1, 11], 1, Concat, [1]],  #未改动前的第10层,在记过上述改动后改为11[[18, 21, 24], 1, Detect, [nc, anchors]],  #17,20,23层改为18,21,24

运行train.py文件可以在输出终端窗口看到上图网络结构,可以看到在第9层已经成功添加EMA注意力机制模块。

                from  n    params  module      arguments                       0                -1  1      3520  models.common.Conv                      [3, 32, 6, 2, 2]                1                -1  1     18560  models.common.Conv                      [32, 64, 3, 2]                  2                -1  1     18816  models.common.C3                        [64, 64, 1]                     3                -1  1     73984  models.common.Conv                      [64, 128, 3, 2]                 4                -1  2    115712  models.common.C3                        [128, 128, 2]                   5                -1  1    295424  models.common.Conv                      [128, 256, 3, 2]                6                -1  3    625152  models.common.C3                        [256, 256, 3]                   7                -1  1   1180672  models.common.Conv                      [256, 512, 3, 2]                8                -1  1   1182720  models.common.C3                        [512, 512, 1]                   9                -1  1     41216  models.EMA.EMA                          [512, 8]                       10                -1  1    656896  models.common.SPPF                      [512, 512, 5]                  11                -1  1    131584  models.common.Conv                      [512, 256, 1, 1]               12                -1  1         0  torch.nn.modules.upsampling.Upsample    [None, 2, 'nearest']           13           [-1, 6]  1         0  models.common.Concat                    [1]14                -1  1    361984  models.common.C3                        [512, 256, 1, False]           15                -1  1     33024  models.common.Conv                      [256, 128, 1, 1]               16                -1  1         0  torch.nn.modules.upsampling.Upsample    [None, 2, 'nearest']           17           [-1, 4]  1         0  models.common.Concat                    [1]18                -1  1     90880  models.common.C3                        [256, 128, 1, False]           19                -1  1    147712  models.common.Conv                      [128, 128, 3, 2]               20          [-1, 15]  1         0  models.common.Concat                    [1]21                -1  1    296448  models.common.C3                        [256, 256, 1, False]           22                -1  1    590336  models.common.Conv                      [256, 256, 3, 2]               23          [-1, 11]  1         0  models.common.Concat                    [1]24                -1  1   1182720  models.common.C3                        [512, 512, 1, False]           25      [18, 21, 24]  1     16182  models.yolo.Detect                      [1, [[10, 13, 16, 30, 33, 23], [30, 61, 62, 45, 59, 119], [116, 90, 156, 198, 373, 326]], [128, 256, 512]]YOLOv5sEMA summary: 222 layers, 7063542 parameters, 7063542 gradients, 16.2 GFLOPs

https://www.bilibili.com/video/BV1s84y1775U/?spm_id_from=333.788&vd_source=f83457e2adc10b543ae4c742fba1e3b2
https://blog.csdn.net/hhhhhhhhhhwwwwwwwwww/article/details/131347981

来源地址:https://blog.csdn.net/qq_43615485/article/details/131470922

阅读原文内容投诉

免责声明:

① 本站未注明“稿件来源”的信息均来自网络整理。其文字、图片和音视频稿件的所属权归原作者所有。本站收集整理出于非商业性的教育和科研之目的,并不意味着本站赞同其观点或证实其内容的真实性。仅作为临时的测试数据,供内部测试之用。本站并未授权任何人以任何方式主动获取本站任何信息。

② 本站未注明“稿件来源”的临时测试数据将在测试完成后最终做删除处理。有问题或投稿请发送至: 邮箱/279061341@qq.com QQ/279061341

软考中级精品资料免费领

  • 历年真题答案解析
  • 备考技巧名师总结
  • 高频考点精准押题
  • 2024年上半年信息系统项目管理师第二批次真题及答案解析(完整版)

    难度     813人已做
    查看
  • 【考后总结】2024年5月26日信息系统项目管理师第2批次考情分析

    难度     354人已做
    查看
  • 【考后总结】2024年5月25日信息系统项目管理师第1批次考情分析

    难度     318人已做
    查看
  • 2024年上半年软考高项第一、二批次真题考点汇总(完整版)

    难度     435人已做
    查看
  • 2024年上半年系统架构设计师考试综合知识真题

    难度     224人已做
    查看

相关文章

发现更多好内容

猜你喜欢

AI推送时光机
位置:首页-资讯-后端开发
咦!没有更多了?去看看其它编程学习网 内容吧
首页课程
资料下载
问答资讯