BAM: Bottleneck Attention Module

BAM: Bottleneck Attention Module

GitHub - Jongchan/attention-module: Official PyTorch code for "BAM: Bottleneck Attention Module (BMVC2018)" and "CBAM: Convolutional Block Attention Module (ECCV2018)"

Given a 3D feature map, BAM produces a 3D attention feature map to emphasize important elements.

We place our module at each bottleneck of models where the downsampling of feature maps occurs.

给定输入特征图，BAM得到一个3D attention map ，经过改进后的特征图通过下式得到

其中表示element-wise mulplication。首先通过两个不同的分支分别计算通道注意力和空间注意力，然后通过下式计算最终的attention map 

其中是sigmoid函数。注意，两个分支的输出需要先resize成，然后再进行相加。

通道分支的计算方法

对于输入特征图，首先是通过全局平均池化得到向量，文中提到："This vector softly encodes global information in each channel "。然后接含一层隐藏层的MLP，即两层全连接层，为了减少额外的参数开销，隐藏层的size设置为，r是reduction ratio，第二个FC再还原回去，这里和SElayer是一样的操作。最后再接一个BN层。

空间分支的计算方法

空间分支得到一个spatial attention map  to emphasize or suppress features in different spatial locations. 具体步骤为：input feature map ，首先经过1×1卷积映射到一个低维空间，这里的r和通道分支的相同；然后经过两层3×3卷积，注意为了增大感受野这里的3×3卷积采用了膨胀卷积dilated convolution；然后再使用1×1卷积映射到；最后再接一个BN层。

合并两个分支的结果

然后需要融合两个分支的结果，在融合之前需要先将两个分支的结果都expand成，这里融合采用的是element-wise summation，然后接sigmoid函数得到最终的attention map。然后将与输入进行element-wise mulplication，再与相加就得到了最终结果refined feature map 。这里借鉴了residual的shortcut结构。

CIFAR-100消融实验

Dilation value and Reduction ratio

论文最终采用dilation value=4, reduction value=16的配置。

Separate or Combined branches

虽然channel和spatial分支都可以提升模型的效果，但结合起来后效果的提升幅度更大。

Combining methods

同样是表(b)中的结果，可以看到，sum的效果最好

Comparison with placing orginal convblocks

作者为了证明BAM带来的效果提升并不是添加了额外的层导致模型更深的作用，因此作者把添加的BAM换成模型原本的block，然后比较两者的效果，从表中结果可以看出，BAM的效果更好。因此得到结论：BAM带来的效果提升并不是因为模型深度的增加，而是BAM本身的结构和注意力机制带来的。

Bottleneck: The efficient point to place BAM

这个实验比较了放置BAM的不同位置，bottlenecks or convolution blocks，结果证明，将BAM放在bottleneck位置可以带来更好的效果并且更少的参数。

官方代码

import torch
import torch.nn as nn
import torch.nn.functional as F


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):
        super(ChannelGate, self).__init__()
        self.gate_c = nn.Sequential()
        self.gate_c.add_module('flatten', Flatten())

        self.gate_c.add_module('gate_c_fc_0', nn.Linear(gate_channel, gate_channel // reduction_ratio))
        self.gate_c.add_module('gate_c_bn_1', nn.BatchNorm1d(gate_channel // reduction_ratio))
        self.gate_c.add_module('gate_c_relu_1', nn.ReLU())
        self.gate_c.add_module('gate_c_fc_final', nn.Linear(gate_channel // reduction_ratio, gate_channel))

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


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, in_tensor):
        return self.gate_s(in_tensor).expand_as(in_tensor)


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, in_tensor):
        att = 1 + F.sigmoid(self.channel_att(in_tensor) * self.spatial_att(in_tensor))
        return att * in_tensor

注意论文中是在每个分支的最终输出加上BN，而在代码中是中间的每一层卷积或是全连接层后都添加BN+ReLU，而最后一层BN和ReLU都不加。