YOLOv5改進之添加CBAM注意力機制的方法

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

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

前言：

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

解決問題：

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

添加方法：

第一步：確定添加的位置，作為即插即用的注意力模塊，可以添加到Y(jié)OLOv5網(wǎng)絡(luò)中的任何地方。本文以添加進卷積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中注冊我們進行修改的Conv_CBAM模塊

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