python

关注公众号 jb51net

关闭
首页 > 脚本专栏 > python > YOLOv5添加CBAM注意力机制

YOLOv5改进之添加CBAM注意力机制的方法

作者:人工智能算法研究院

注意力机制最先被用在NLP领域,Attention就是为了让模型认识到数据中哪一部分是最重要的,为它分配更大的权重,获得更多的注意力在一些特征上,让模型表现更好,这篇文章主要给大家介绍了关于YOLOv5改进之添加CBAM注意力机制的相关资料,需要的朋友可以参考下

前 言:

作为当前先进的深度学习目标检测算法YOLOv5,已经集合了大量的trick,但是在处理一些复杂背景问题的时候,还是容易出现错漏检的问题。此后的系列文章,将重点对YOLOv5的如何改进进行详细的介绍,目的是为了给那些搞科研的同学需要创新点或者搞工程项目的朋友需要达到更好的效果提供自己的微薄帮助和参考。

解决问题:

加入CBAM双通道注意力机制,可以让网络更加关注待检测目标,提高检测效果

添加方法:

第一步:确定添加的位置,作为即插即用的注意力模块,可以添加到YOLOv5网络中的任何地方。本文以添加进卷积Conv模块中为例。

第二步:common.py构建融入CBAM模块的Conv_CBAM,与原Conv模块不同的是:在该模块中的激活函数后加入CBAM模块。

class Conv_CBAM(nn.Module):
    # Standard convolution
    def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True):  # ch_in, ch_out, kernel, stride, padding, groups
        super(Conv_CBAM, self).__init__()
        self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=False)
        self.bn = nn.BatchNorm2d(c2)
        self.act = nn.Hardswish() if act else nn.Identity()
        self.ca = ChannelAttention(c2)
        self.sa = SpatialAttention()
 
    def forward(self, x):
        x = self.act(self.bn(self.conv(x)))
        x = self.ca(x) * x
        x = self.sa(x) * x
        return x
 
    def fuseforward(self, x):
        return self.act(self.conv(x))

第三步:yolo.py中注册我们进行修改的Conv_CBAM模块

第四步:修改yaml文件,本文以修改主干特征提取网络为例,将原Conv模块改为Conv_CBAM即可。

第五步:将train.py中改为本文的yaml文件即可,开始训练。

结 果:

本人在多个数据集上做了大量实验,针对不同的数据集效果不同,同一个数据集的不同添加位置方法也是有差异,需要大家进行实验。有效果有提升的情况占大多数。

PS:CBAM双通道注意力机制,不仅仅是可以添加进YOLOv5,也可以添加进任何其他的深度学习网络,不管是分类还是检测还是分割,主要是计算机视觉领域,都可能会有不同程度的提升效果。

总 结

到此这篇关于YOLOv5改进之添加CBAM注意力机制的文章就介绍到这了,更多相关YOLOv5添加CBAM注意力机制内容请搜索脚本之家以前的文章或继续浏览下面的相关文章希望大家以后多多支持脚本之家!

您可能感兴趣的文章:
阅读全文