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模塊有一點相似,但是要更豐富一點。
-
最終按照輸入tensor x的形狀進行擴展,得到關於通道的注意力。
然后講一下與CBAM中的區別:
- CBAM中使用的是卷積實現的通道處理,這里使用的是線性變換Linear,但是背后的數學原理應該是一致的,區別在於計算量上
- 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進行歸一化。
后記:感覺BAM跟CBAM相比有一點點復雜,沒有CBAM的那種簡潔美。這兩篇都是坐着在同一時期進行發表的,所以並沒有互相的一個詳細的對照,但是大概看了一下,感覺CBAM效果好於BAM。