【CV中的Attention機制】并聯(lián)版的CBAM-BAM模塊

前言:之前介紹了CBAM模塊,可以方便的添加到自己的網(wǎng)絡(luò)模型中,代碼比較簡單容易理解。CBAM模塊的實現(xiàn)是通過先后施加通道注意力和空間注意力完成信息的提煉。今天介紹的這篇文章也是來自CBAM團隊,可以理解為空間注意力機制和通道注意力機制的并聯(lián),但是具體實現(xiàn)與CBAM有較大差別,雖然代碼量相對而言比較大,實際表達的內(nèi)容并不復(fù)雜。

  • 作者:pprp
  • 編輯:BBuf

1. BAM

BAM全程是bottlenect attention module,與CBAM很相似的起名,還是CBAM的團隊完成的作品。

CBAM被ECCV18接收,BAM被BMVC18接收。

CBAM可以看做是通道注意力機制和空間注意力機制的串聯(lián)(先通道后空間),BAM可以看做兩者的并聯(lián)。

image

這個模塊之所以叫bottlenect是因為這個模塊放在DownSample 也就是pooling layer之前,如下圖所示:

image

由于改論文與上一篇:CBAM模塊的理論部分極為相似,下邊直接進行算法實現(xiàn)部分。

2. 通道部分的實現(xiàn)

class Flatten(nn.Module):
    def forward(self, x):
        return x.view(x.size(0), -1)
    
class ChannelGate(nn.Module):
    def __init__(self, gate_channel, reduction_ratio=16, num_layers=1):
        super(ChannelGate, self).__init__()
        self.gate_c = nn.Sequential()
        self.gate_c.add_module('flatten', Flatten())

        gate_channels = [gate_channel]  # eg 64
        gate_channels += [gate_channel // reduction_ratio] * num_layers  # eg 4
        gate_channels += [gate_channel]  # 64
        # gate_channels: [64, 4, 4]

        for i in range(len(gate_channels) - 2):
            self.gate_c.add_module(
                'gate_c_fc_%d' % i,
                nn.Linear(gate_channels[i], gate_channels[i + 1]))
            self.gate_c.add_module('gate_c_bn_%d' % (i + 1),
                                   nn.BatchNorm1d(gate_channels[i + 1]))
            self.gate_c.add_module('gate_c_relu_%d' % (i + 1), nn.ReLU())

        self.gate_c.add_module('gate_c_fc_final',
                               nn.Linear(gate_channels[-2], gate_channels[-1]))

    def forward(self, x):
        avg_pool = F.avg_pool2d(x, x.size(2), stride=x.size(2))
        return self.gate_c(avg_pool).unsqueeze(2).unsqueeze(3).expand_as(x)

看上去代碼要比CBAM中的ChannelAttention模塊要多很多,貼上ChannelAttention代碼方便對比:

class ChannelAttention(nn.Module):
    def __init__(self, in_planes, rotio=16):
        super(ChannelAttention, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.max_pool = nn.AdaptiveMaxPool2d(1)

        self.sharedMLP = nn.Sequential(
            nn.Conv2d(in_planes, in_planes // ratio, 1, bias=False), nn.ReLU(),
            nn.Conv2d(in_planes // rotio, in_planes, 1, bias=False))
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        avgout = self.sharedMLP(self.avg_pool(x))
        maxout = self.sharedMLP(self.max_pool(x))
        return self.sigmoid(avgout + maxout)

首先講ChannelGate的處理流程:

  • 使用avg_pool2d測試

    >>> import torch.nn.functional as F
    >>> import torch
    >>> x = torch.ones((12, 8, 64, 64))
    >>> x.shape
    torch.Size([12, 8, 64, 64])
    >>> F.avg_pool2d(x,x.size(2), stride=x.size(2)).shape
    torch.Size([12, 8, 1, 1])
    >>>
    

    其效果與AdaptiveAvgPool2d(1)是一樣的。

  • 然后經(jīng)過gate_c模塊,里邊先經(jīng)過Flatten將其變?yōu)閇batch size, channel]形狀的tensor, 然后后邊一大部分都是Linear模塊,進行線性變換。(ps:雖然代碼看上去多,但是功能很簡單)這個部分與SE模塊有一點相似,但是可以添加多個Linear層,蘊含的信息要更豐富一點。

  • 最終按照輸入tensor x的形狀進行擴展,得到關(guān)于通道的注意力。

然后講一下與CBAM中的channel attention的區(qū)別:

  • CBAM中使用的是先用adaptiveAvgPooling,然后進行卷積實現(xiàn)的通道處理;BAM使用的也是adaptiveAvgPooling, 然后進行多個Linear線性變換,得到channel attention。其實關(guān)于用1\times1卷積和Linear層實現(xiàn),在feature map尺寸為1\times?1的時候,兩者從數(shù)學(xué)原理上講,沒有區(qū)別。具體可以參考知乎上的問題:1*1的卷積核和全連接層有什么異同?
  • CBAM中激活函數(shù)使用sigmoid, BAM中的通道部分使用了ReLU,還添加了BN層。

3. 空間注意力機制

class SpatialGate(nn.Module):
    def __init__(self,
                 gate_channel,
                 reduction_ratio=16,
                 dilation_conv_num=2,
                 dilation_val=4):
        super(SpatialGate, self).__init__()
        self.gate_s = nn.Sequential()

        self.gate_s.add_module(
            'gate_s_conv_reduce0',
            nn.Conv2d(gate_channel,
                      gate_channel // reduction_ratio,
                      kernel_size=1))
        self.gate_s.add_module('gate_s_bn_reduce0',
                               nn.BatchNorm2d(gate_channel // reduction_ratio))
        self.gate_s.add_module('gate_s_relu_reduce0', nn.ReLU())

        # 進行多個空洞卷積,豐富感受野
        for i in range(dilation_conv_num):
            self.gate_s.add_module(
                'gate_s_conv_di_%d' % i,
                nn.Conv2d(gate_channel // reduction_ratio,
                          gate_channel // reduction_ratio,
                          kernel_size=3,
                          padding=dilation_val,
                          dilation=dilation_val))
            self.gate_s.add_module(
                'gate_s_bn_di_%d' % i,
                nn.BatchNorm2d(gate_channel // reduction_ratio))
            self.gate_s.add_module('gate_s_relu_di_%d' % i, nn.ReLU())

        self.gate_s.add_module(
            'gate_s_conv_final',
            nn.Conv2d(gate_channel // reduction_ratio, 1, kernel_size=1))

    def forward(self, x):
        return self.gate_s(x).expand_as(x)

這里可以看出,代碼量相比CBAM中的spatial attention要大很多,依然進行對比:

class SpatialAttention(nn.Module):
    def __init__(self, kernel_size=7):
        super(SpatialAttention, self).__init__()
        assert kernel_size in (3,7), "kernel size must be 3 or 7"
        padding = 3 if kernel_size == 7 else 1

        self.conv = nn.Conv2d(2,1,kernel_size, padding=padding, bias=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        avgout = torch.mean(x, dim=1, keepdim=True)
        maxout, _ = torch.max(x, dim=1, keepdim=True)
        x = torch.cat([avgout, maxout], dim=1)
        x = self.conv(x)
        return self.sigmoid(x)

這個部分空間注意力處理就各有特色了,先說一下BAM中的流程:

  • 先經(jīng)過一個conv+bn+relu模塊,通道縮進,信息進行壓縮。
  • 然后經(jīng)過了多個dilated conv+bn+relu模塊,空洞率設(shè)置為4(默認)。
  • 最后經(jīng)過一個卷積,將通道壓縮到1。
  • 最終將其擴展為tensor x的形狀。

區(qū)別在于:

  • CBAM中通過通道間的max,avg處理成通道數(shù)為2的feature, 然后通過卷積+Sigmoid得到最終的map
  • BAM中則全部通過卷積或者空洞卷積完成信息處理,計算量更大一點, 但是融合了多感受野,信息更加豐富。

4. BAM融合

class BAM(nn.Module):
    def __init__(self, gate_channel):
        super(BAM, self).__init__()
        self.channel_att = ChannelGate(gate_channel)
        self.spatial_att = SpatialGate(gate_channel)

    def forward(self, x):
        att = 1 + F.sigmoid(self.channel_att(x) * self.spatial_att(x))
        return att * x

最終融合很簡單,需要注意的就是兩者是相乘的,并且使用了sigmoid進行歸一化。


論文鏈接:https://arxiv.org/pdf/1807.06514

核心代碼:https://github.com/pprp/SimpleCVReproduction/tree/master/attention/BAM

后記:感覺BAM跟CBAM相比有一點點復(fù)雜,沒有CBAM的那種簡潔美。這兩篇都是坐著在同一時期進行發(fā)表的,所以并沒有互相的一個詳細的對照,但是大概看了一下,感覺CBAM效果好于BAM。

?著作權(quán)歸作者所有,轉(zhuǎn)載或內(nèi)容合作請聯(lián)系作者
【社區(qū)內(nèi)容提示】社區(qū)部分內(nèi)容疑似由AI輔助生成,瀏覽時請結(jié)合常識與多方信息審慎甄別。
平臺聲明:文章內(nèi)容(如有圖片或視頻亦包括在內(nèi))由作者上傳并發(fā)布,文章內(nèi)容僅代表作者本人觀點,簡書系信息發(fā)布平臺,僅提供信息存儲服務(wù)。

相關(guān)閱讀更多精彩內(nèi)容

友情鏈接更多精彩內(nèi)容