YOLOv5改進教程之添加註意力機制

本文主要給大傢講解一下,如何在yolov5中添加註意力機制,

這裡提供SE通道註意力的改進方法,其他註意力的添加方法,大同小異

首先找到SE註意力機制的pytorch代碼

class SELayer(nn.Module):
    def __init__(self, c1, r=16):
        super(SELayer, self).__init__()
        self.avgpool = nn.AdaptiveAvgPool2d(1)
        self.l1 = nn.Linear(c1, c1 // r, bias=False)
        self.relu = nn.ReLU(inplace=True)
        self.l2 = nn.Linear(c1 // r, c1, bias=False)
        self.sig = nn.Sigmoid()
 
    def forward(self, x):
        b, c, _, _ = x.size()
        y = self.avgpool(x).view(b, c)
        y = self.l1(y)
        y = self.relu(y)
        y = self.l2(y)
        y = self.sig(y)
        y = y.view(b, c, 1, 1)
        return x * y.expand_as(x)

直接把SE註意力機制的程序,復制到models文件夾下面的common.py文件中

然後找到yolo.py,在這個文件中找到下面這一行

 然後把SE添加到這個註冊表裡,

 直接在C3Ghost後面加上SELayer

然後是要修改yaml文件,對於SE這種即插即用的註意力機制,輸入和輸出的通道數相同,所以不會影響其他模塊的運行

註意力機制也可以插在其他地方,隻要註意把通道數對應好就行

然後head部分也要進行相應的修改,因為我們多加瞭一層,所以需要保持head以及最後的輸出層不變就得稍微改一下head部分,同樣,註意力機制也可以放在head裡面,跟加在backbone裡面的方法相同。

這是原始的head部分,需要修改成下面這樣

因為我們把SE註意力機制插在第八層之後,所以原先的14層就變成瞭15層,同樣的,最後的輸出也需要把層數加一。

添加這些註意力機制是yolov5最基礎的改進,但是說實話,註意力機制可以添加在很多的地方,不一定會有效果,所以插在哪裡效果最好就需要大傢自行討論瞭。 註意力機制也有很多種,se註意力機制應該是屬於最基礎的通道註意力瞭吧

總結

到此這篇關於YOLOv5改進教程之添加註意力機制的文章就介紹到這瞭,更多相關YOLOv5添加註意力機制內容請搜索WalkonNet以前的文章或繼續瀏覽下面的相關文章希望大傢以後多多支持WalkonNet!

推薦閱讀: