【注意力机制】CBAM详解（文末附代码）

论文题目：《CBAM: Convolutional Block Attention Module》
论文地址：https://arxiv.org/pdf/1807.06521.pdf

1. 前言

       论文（2018年）提出了一种轻量的注意力模块( CBAM，Convolutional Block Attention Module )，可以在通道和空间维度上进行 Attention 。论文在 ResNet 和 MobileNet 等经典结构上添加了 CBAM 模块并进行对比分析，同时也进行了可视化，发现 CBAM 更关注识别目标物体，这也使得 CBAM 具有更好的解释性。

2. CBAM（Convolutional Block Attention Module）

       CBAM结构如下图所示：

       可以看到 CBAM 包含2个独立的子模块， 通道注意力模块（Channel Attention Module，CAM) 和空间注意力模块（Spartial Attention Module，SAM) ，分别进行通道与空间上的 Attention 。 这样不只能够节约参数和计算力，并且保证了其能够做为即插即用的模块集成到现有的网络架构中去。通道上的 Attention 机制在 2017 年的 SENet 就被提出，SENet可以参考我的这篇文章。事实上，CAM 与 SENet 相比，只是多了一个并行的 Max Pooling 层。至于为何如此更改，论文也给出了解释和实验数据支持。

2.1. Channel Attention Module（CAM）

       通道上的Attention模块以及具体计算如下图所示：

具体流程如下：
       将输入的特征图F（H×W×C）分别经过基于width和height的global max pooling（全局最大池化）和global average pooling（全局平均池化），得到两个1×1×C的特征图，接着，再将它们分别送入一个两层的神经网络（MLP），第一层神经元个数为 C/r（r为减少率），激活函数为 Relu，第二层神经元个数为 C，这个两层的神经网络是共享的。而后，将MLP输出的特征进行基于element-wise的加和操作，再经过sigmoid激活操作，生成最终的channel attention feature，即M_c。最后，将M_c和输入特征图F做element-wise乘法操作，生成Spatial attention模块需要的输入特征。

● pooling的使用

       在channel attention中，表1对于pooling的使用进行了实验对比，发现avg & max的并行池化的效果要更好。这里也有可能是池化丢失的信息太多，avg&max的并行连接方式比单一的池化丢失的信息更少，所以效果会更好一点。

2.2. Spatial Attention Module（SAM）

       空间上的Attention模块以及具体计算如下图所示：

具体流程如下：
       将Channel attention模块输出的特征图F‘作为本模块的输入特征图。首先做一个基于channel的global max pooling 和global average pooling，得到两个H×W×1 的特征图，然后将这2个特征图基于channel 做concat操作（通道拼接）。然后经过一个7×7卷积（7×7比3×3效果要好）操作，降维为1个channel，即H×W×1。再经过sigmoid生成spatial attention feature，即M_s。最后将该feature和该模块的输入feature做乘法，得到最终生成的特征。

3. CAM和SAM的组合形式

       通道注意力和空间注意力这两个模块能够以并行或者串行顺序的方式组合在一块儿，关于通道和空间上的串行顺序和并行作者进行了实验对比，发现先通道再空间的结果会稍微好一点。具体实验结果如下：

       从表3中可以看出，在ResNet50的基准架构尚，两个attetnion子模块的连接顺序里面的确是channel + spatial的要更好一些，也要好于标准的SENet的通道attention。

● CBAM和ResBlock组合

4. CBAM可视化

       利用 Grad-CAM 对不一样的网络进行可视化后，能够发现，引入 CBAM 后，特征覆盖到了待识别物体的更多部位，而且最终判别物体的几率也更高，这代表注意力机制的确让网络学会了关注重点信息。

5. 总结

       论文提出了一种基于注意力机制的轻量型结构 CBAM ，基本上可以添加到所有常规的卷积层中。
       文中验证了 Channel Attention Module 中 avg 与 max 并行的方式最好，接下来通过实验验证了 Channel Attention Module 和 Spatial Attention Module 的最佳先后顺序，还对比了 CBAM 与 SENet 的性能。
       文章还可视化了 CBAM 的关注区域，使得 CBAM 具有更好的解释性。最后在目标检测任务上进行实验，验证了 CBAM 的通用性好。

代码：

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

class BasicConv(nn.Module):
    def __init__(self, in_planes, out_planes, kernel_size, stride=1, padding=0, dilation=1, groups=1, relu=True, bn=True, bias=False):
        super(BasicConv, self).__init__()
        self.out_channels = out_planes
        self.conv = nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias)
        self.bn = nn.BatchNorm2d(out_planes,eps=1e-5, momentum=0.01, affine=True) if bn else None
        self.relu = nn.ReLU() if relu else None

    def forward(self, x):
        x = self.conv(x)
        if self.bn is not None:
            x = self.bn(x)
        if self.relu is not None:
            x = self.relu(x)
        return x

class Flatten(nn.Module):
    def forward(self, x):
        return x.view(x.size(0), -1)

class ChannelGate(nn.Module):
    def __init__(self, gate_channels, reduction_ratio=16, pool_types=['avg', 'max']):
        super(ChannelGate, self).__init__()
        self.gate_channels = gate_channels
        self.mlp = nn.Sequential(
            Flatten(),
            nn.Linear(gate_channels, gate_channels // reduction_ratio),
            nn.ReLU(),
            nn.Linear(gate_channels // reduction_ratio, gate_channels)
            )
        self.pool_types = pool_types
    def forward(self, x):
        channel_att_sum = None
        for pool_type in self.pool_types:
            if pool_type=='avg':
                avg_pool = F.avg_pool2d( x, (x.size(2), x.size(3)), stride=(x.size(2), x.size(3)))
                channel_att_raw = self.mlp( avg_pool )
            elif pool_type=='max':
                max_pool = F.max_pool2d( x, (x.size(2), x.size(3)), stride=(x.size(2), x.size(3)))
                channel_att_raw = self.mlp( max_pool )
            elif pool_type=='lp':
                lp_pool = F.lp_pool2d( x, 2, (x.size(2), x.size(3)), stride=(x.size(2), x.size(3)))
                channel_att_raw = self.mlp( lp_pool )
            elif pool_type=='lse':
                # LSE pool only
                lse_pool = logsumexp_2d(x)
                channel_att_raw = self.mlp( lse_pool )

            if channel_att_sum is None:
                channel_att_sum = channel_att_raw
            else:
                channel_att_sum = channel_att_sum + channel_att_raw

        scale = F.sigmoid( channel_att_sum ).unsqueeze(2).unsqueeze(3).expand_as(x)
        return x * scale

def logsumexp_2d(tensor):
    tensor_flatten = tensor.view(tensor.size(0), tensor.size(1), -1)
    s, _ = torch.max(tensor_flatten, dim=2, keepdim=True)
    outputs = s + (tensor_flatten - s).exp().sum(dim=2, keepdim=True).log()
    return outputs

class ChannelPool(nn.Module):
    def forward(self, x):
        return torch.cat( (torch.max(x,1)[0].unsqueeze(1), torch.mean(x,1).unsqueeze(1)), dim=1 )

class SpatialGate(nn.Module):
    def __init__(self):
        super(SpatialGate, self).__init__()
        kernel_size = 7
        self.compress = ChannelPool()
        self.spatial = BasicConv(2, 1, kernel_size, stride=1, padding=(kernel_size-1) // 2, relu=False)
    def forward(self, x):
        x_compress = self.compress(x)
        x_out = self.spatial(x_compress)
        scale = F.sigmoid(x_out) # broadcasting
        return x * scale

class CBAM(nn.Module):
    def __init__(self, gate_channels, reduction_ratio=16, pool_types=['avg', 'max'], no_spatial=False):
        super(CBAM, self).__init__()
        self.ChannelGate = ChannelGate(gate_channels, reduction_ratio, pool_types)
        self.no_spatial=no_spatial
        if not no_spatial:
            self.SpatialGate = SpatialGate()
    def forward(self, x):
        x_out = self.ChannelGate(x)
        if not self.no_spatial:
            x_out = self.SpatialGate(x_out)
        return x_out