【CV中的Attention機制】並行版的CBAM-BAM模組

  • 2020 年 2 月 12 日
  • 筆記

前言：之前介紹了CBAM模組，可以方便的添加到自己的網路模型中，程式碼比較簡單容易理解。CBAM模組的實現是通過先後施加通道注意力和空間注意力完成資訊的提煉。今天介紹的這篇文章也是來自CBAM團隊，可以理解為空間注意力機制和通道注意力機制的並聯，但是具體實現與CBAM有較大差別，雖然程式碼量相對而言比較大，實際表達的內容並不複雜。

  • 作者：pprp
  • 編輯：BBuf

1. BAM

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

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

CBAM可以看做是通道注意力機制和空間注意力機制的串聯（先通道後空間），BAM可以看做兩者的並聯。

這個模組之所以叫bottlenect是因為這個模組放在DownSample 也就是pooling layer之前，如下圖所示：

由於改論文與上一篇：CBAM模組的理論部分極為相似，下邊直接進行演算法實現部分。

2. 通道部分的實現

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)是一樣的。
  • 然後經過gate_c模組，裡邊先經過Flatten將其變為[batch size, channel]形狀的tensor, 然後後邊一大部分都是Linear模組，進行線性變換。（ps:雖然程式碼看上去多，但是功能很簡單）這個部分與SE模組有一點相似，但是可以添加多個Linear層，蘊含的資訊要更豐富一點。
  • 最終按照輸入tensor x的形狀進行擴展，得到關於通道的注意力。

然後講一下與CBAM中的channel attention的區別：

  • CBAM中使用的是先用adaptiveAvgPooling，然後進行卷積實現的通道處理；BAM使用的也是adaptiveAvgPooling, 然後進行多個Linear線性變換，得到channel attention。其實關於用11卷積和Linear層實現，在feature map尺寸為11的時候，兩者從數學原理上講，沒有區別。具體可以參考知乎上的問題：1×1的卷積核和全連接層有什麼異同？
  • CBAM中激活函數使用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中的流程：

  • 先經過一個conv+bn+relu模組，通道縮進，資訊進行壓縮。
  • 然後經過了多個dilated conv+bn+relu模組，空洞率設置為4（默認）。
  • 最後經過一個卷積，將通道壓縮到1。
  • 最終將其擴展為tensor x的形狀。

區別在於：

  • CBAM中通過通道間的max,avg處理成通道數為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

往期文章鏈接-CBAM模組

後記：BAM跟CBAM從程式碼層面相比有一點點複雜，沒有CBAM的那種簡潔美。這兩篇都是坐著在同一時期進行發表的，所以並沒有互相的一個詳細的對照，但是大體來看，CBAM效果好於BAM。

VirMach 便宜 VPS

QNews

QNews