BAM: Bottleneck Attention Module__实现

BAM: Bottleneck Attention Module

引言

  在此论文中，我们把重心放在了Attention对于一般深度神经网络的影响上， 我们提出了一个简单但是有效的Attention 模型—BAM，它可以结合到任何前向传播卷积神经网络中，我们的模型通过两个分离的路径 channel和spatial, 得到一个Attention Map.
整体结构图如下：

这里作者将BAM放在了Resnet网络中每个stage之间。有趣的是，通过可视化我们可以看到多层BAMs形成了一个分层的注意力机制，这有点像人类的感知机制。BAM在每个stage之间消除了像背景语义特征这样的低层次特征，然后逐渐聚焦于高级的语义–明确的目标（比如图中的狗）.

主要思想：

channel attention branch

对于给定的feature map$\mathrm{F}\in R^{C\ast H\ast W}$,BAM可以得到一个3D的Attention map $M(F)\in R^{C\ast H\ast W}$.加强的feature map F′;
${\mathrm{F}}^{\prime }=\mathrm{F}+\mathrm{F}\otimes \mathrm{M}(\mathrm{F})$

为了设计一个有效且强大的模型，我们首先计算channel attention,然后计算spatial attention.这时M(F)就变成：
$\mathrm{M}(\mathrm{F})=\sigma \left({\mathrm{M}}_c(\mathrm{F})+{\mathrm{M}}_s(\mathrm{F})\right)$

这里σ 代表sigmoid函数，为了聚合feature map在每个通道维度，我们采用全局平均池化得到$F_C$这个向量然后对全局信息在每个通道进行软编码。为了评估Attention在每个通道的效果？我们使用了一个多层感知(MLP)用一层隐藏层。在MLP之后，我们增加了BN去调整规模和空间分支一样的输出，channel attention可以被计算为：

${\mathbf{M}}_{\mathbf{c}}(\mathbf{F})=BN(MLP(\text{AvgPool}(\mathbf{F})))$
$=BN\left({\mathbf{W}}_1\left({\mathbf{W}}_{0}AvgPool(\mathbf{F})+{\mathbf{b}}_0\right)+{\mathbf{b}}_1\right)$
where ${\mathbf{W}}_0\in {\mathbb{R}}^{C/r\times C},{\mathbf{b}}_0\in {\mathbb{R}}^{C/r},{\mathbf{W}}_1\in {\mathbb{R}}^{C\times C/r},{\mathbf{b}}_1\in {\mathbb{R}}^C$

Spatial attention branch

这个空间分支产生了空间Attention去增强或者抑制特征在不同的空间位置，众所周知，利用上下文信息是去知道应该关注哪些位置的关键点。在这里我们为了高效性运用空洞卷积去增大感受野。
我们观察到，与标准卷积相比，空洞卷积有助于构造更有效的spatial map.
细节图：

空洞模型结构 给与中间feature map F,这个module 计算对应的Attention mapM(F)通过两个单独的Attention 分支–channle Mc 和空间${\mathrm{M}}_S$.这里有两个超参数 dilation value (d)和reduction ratio®. d参数决定了感受野大小,这对空间分支聚合上下文信息非常重要。这里我们set d=4 r=16.

我们采用空洞卷积来高效扩大感受野。我们观察到空洞卷积有助于构建比标准卷积更有效的空间映射。 我们的空间分支采用了ResNet建议的“瓶颈结构”，既节省了参数数量又节省了计算开销。 具体地，使用1×1卷积将特征$\mathbf{F}\in {\mathbb{R}}^{C\times H\times W}$投影到缩小尺寸的${\mathbb{R}}^{C/r\times H\times W}$，以在整个通道维度上对特征图进行结合和压缩。 为简单起见，我们使用与通道分支相同的缩减比r。 在减少之后，应用两个3×3扩张卷积以有效地利用上下文信息。 最后，使用1×1卷积将特征再次简化为${\mathbb{R}}^{1\times H\times W}$空间注意力图。 对于缩放调整，在空间分支的末尾应用批量标准化层。 简而言之，空间注意力计算如下：
${\mathbf{M}}_{\mathbf{s}}(\mathbf{F})=BN\left(f_3^{1\times 1}\left(f_2^{3\times 3}\left(f_1^{3\times 3}\left(f_0^{1\times 1}(\mathbf{F})\right)\right)\right)\right)$
其中f表示卷积运算，BN表示批量归一化运算，上标表示卷积滤波器大小。 通道缩减有两个1×1卷积。中间3×3扩张卷积用于聚合具有较大感受野的上下文信息。

Combine two attention branches

在从两个注意力分支中获取通道注意力Mc(F)和空间注意力${\mathrm{M}}_S$(F)后，我们将它们组合起来，生成最终的3D注意力mapM(F)。由于这两个注意图的形状不同，我们将注意图扩展到$ R^{CHW}$，然后将它们合并。在逐项求和、乘法、max运算等多种组合方法中，针对梯度流的特点,选择逐项求和。我们通过实证验证了基于元素的求和在三种选择中效果最好。求和后，我们取一个sigmoid函数，得到0到1范围内的最终三维注意映射M(F)。将该三维注意图与输入特征图F巧妙相乘，然后将其添加到原始输入特征图上，得到细化后的特征图F′
${\mathrm{F}}^{\prime }=\mathrm{F}+\mathrm{F}\otimes \mathrm{M}(\mathrm{F})$

Pytorch实现BAM

源码github链接

import torch
import math
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, num_layers=1):
        super(ChannelGate, self).__init__()
        self.gate_c = nn.Sequential()
        # after avg_pool
        self.gate_c.add_module('flatten', Flatten())
        gate_channels = [gate_channel]
        gate_channels += [gate_channel // reduction_ratio] * num_layers
        gate_channels += [gate_channel]
        for i in range(len(gate_channels) - 2):
            # fc->bn
            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.BatchNorm2d(gate_channels[i+1]))
            self.gate_c.add_module('gate_c_relu_%d'%(i+1), nn.ReLU())
        # final_fc
        self.gate_c.add_module('gate_c_fc_final', nn.Linear(gate_channels[-2], gate_channels[-1]))

    def forward(self, in_tensor):
        # Global avg pool
        avg_pool = F.avg_pool2d(in_tensor, in_tensor.size(2), stride=in_tensor.size(2))
        # C∗H∗W -> C*1*1 -> C*H*W
        return self.gate_c(avg_pool).unsqueeze(2).unsqueeze(3).expand_as(in_tensor)

class SpatiaGate(nn.Module):
    # dilation value and reduction ratio, set d = 4 r = 16
    def __init__(self, gate_channel, reduction_ratio=16, dilation_conv_num=2, dilation_val=4):
        self.gate_s = nn.Sequential()
        # 1x1 + (3x3)*2 + 1x1
        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))  # 1×H×W

    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 = SpatiaGate(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

将BAM融入到Resnet中，在下一篇CBAM中展示源码。