BAM: Bottleneck Attention Module

引言

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

整体结构图如下：

在这里插入图片描述

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

主要思想：

channel attention branch

对于给定的feature map,BAM可以得到一个3D的Attention map .加强的feature map F′;

为了设计一个有效且强大的模型，我们首先计算channel attention,然后计算spatial attention.这时M(F)就变成：

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

where

Spatial attention branch

这个空间分支产生了空间Attention去增强或者抑制特征在不同的空间位置，众所周知，利用上下文信息是去知道应该关注哪些位置的关键点。在这里我们为了高效性运用空洞卷积去增大感受野。

我们观察到，与标准卷积相比，空洞卷积有助于构造更有效的spatial map.

细节图：

在这里插入图片描述

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

我们采用空洞卷积来高效扩大感受野。我们观察到空洞卷积有助于构建比标准卷积更有效的空间映射。 我们的空间分支采用了ResNet建议的“瓶颈结构”，既节省了参数数量又节省了计算开销。 具体地，使用1×1卷积将特征投影到缩小尺寸的，以在整个通道维度上对特征图进行结合和压缩。 为简单起见，我们使用与通道分支相同的缩减比r。 在减少之后，应用两个3×3扩张卷积以有效地利用上下文信息。 最后，使用1×1卷积将特征再次简化为空间注意力图。 对于缩放调整，在空间分支的末尾应用批量标准化层。 简而言之，空间注意力计算如下：

其中f表示卷积运算，BN表示批量归一化运算，上标表示卷积滤波器大小。 通道缩减有两个1×1卷积。中间3×3扩张卷积用于聚合具有较大感受野的上下文信息。

Combine two attention branches

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