BAM: Bottleneck Attention Module

2022-01-23 23:32:05 阅读：327 来源： 互联网

标签：Bottleneck nn self Attention module gate BAM channel

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.

给定输入特征图 $\small F\in \mathbb{R}^{C\times H\times W}$ ，BAM得到一个3D attention map $\small M(F)\in \mathbb{R}^{C\times H\times W}$ ，经过改进后的特征图 $\small F^{'}$ 通过下式得到

其中 $\small \bigotimes$ 表示element-wise mulplication。首先通过两个不同的分支分别计算通道注意力 $\small M_{c}(F)\in \mathbb{R}^{C}$ 和空间注意力 $\small M_{s}(F)\in \mathbb{R}^{H\times W}$ ，然后通过下式计算最终的attention map $\small M(F)$

其中 $\small \sigma$ 是sigmoid函数。注意，两个分支的输出需要先resize成 $\small \mathbb{R}^{C\times H\times W}$ ，然后再进行相加。

通道分支的计算方法

$\small F\in \mathbb{R}^{C\times H\times W}$

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

空间分支的计算方法

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

合并两个分支的结果

然后需要融合两个分支的结果，在融合之前需要先将两个分支的结果都expand成 $\small \mathbb{R}^{C\times H\times W}$ ，这里融合采用的是element-wise summation，然后接sigmoid函数得到最终的attention map $\small M(F)\in \mathbb{R}^{C\times H\times W}$ 。然后将 $\small M(F)$ 与输入 $\small F$ 进行element-wise mulplication，再与 $\small F$ 相加就得到了最终结果refined feature map $\small F^{'}$ 。这里借鉴了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都不加。

标签：Bottleneck,nn,self,Attention,module,gate,BAM,channel
来源： https://blog.csdn.net/ooooocj/article/details/122599779