推荐|BAM: Bottleneck Attention Module_

文章目录

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 * H * W}$ ,BAM可以得到一个3D的Attention map $M(F)∈R^{C*H*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})=B N(M L P(\text {AvgPool}(\mathbf{F})))$
$N\left(\mathbf{W}_{1}\left(\mathbf{W}_{0} A v g P o o l(\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})=B N\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中展示源码。

BAM: Bottleneck Attention Module__实现