快捷導(dǎo)航

YOLOv5改進(jìn)之添加CBAM注意力機(jī)制的方法

更新時(shí)間：2022年11月18日 11:09:07 作者：人工智能算法研究院

注意力機(jī)制最先被用在NLP領(lǐng)域,Attention就是為了讓模型認(rèn)識(shí)到數(shù)據(jù)中哪一部分是最重要的,為它分配更大的權(quán)重,獲得更多的注意力在一些特征上,讓模型表現(xiàn)更好,這篇文章主要給大家介紹了關(guān)于YOLOv5改進(jìn)之添加CBAM注意力機(jī)制的相關(guān)資料,需要的朋友可以參考下

前言：

作為當(dāng)前先進(jìn)的深度學(xué)習(xí)目標(biāo)檢測(cè)算法YOLOv5，已經(jīng)集合了大量的trick，但是在處理一些復(fù)雜背景問題的時(shí)候，還是容易出現(xiàn)錯(cuò)漏檢的問題。此后的系列文章，將重點(diǎn)對(duì)YOLOv5的如何改進(jìn)進(jìn)行詳細(xì)的介紹，目的是為了給那些搞科研的同學(xué)需要?jiǎng)?chuàng)新點(diǎn)或者搞工程項(xiàng)目的朋友需要達(dá)到更好的效果提供自己的微薄幫助和參考。

解決問題：

加入CBAM雙通道注意力機(jī)制，可以讓網(wǎng)絡(luò)更加關(guān)注待檢測(cè)目標(biāo)，提高檢測(cè)效果

添加方法：

第一步：確定添加的位置，作為即插即用的注意力模塊，可以添加到Y(jié)OLOv5網(wǎng)絡(luò)中的任何地方。本文以添加進(jìn)卷積Conv模塊中為例。

第二步：common.py構(gòu)建融入CBAM模塊的Conv_CBAM,與原Conv模塊不同的是：在該模塊中的激活函數(shù)后加入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中注冊(cè)我們進(jìn)行修改的Conv_CBAM模塊

第四步：修改yaml文件，本文以修改主干特征提取網(wǎng)絡(luò)為例，將原Conv模塊改為Conv_CBAM即可。