CBAM重点干货和流程详解及Pytorch实现

---

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

前言

CBAM是轻量级的卷积块注意力模型，它可以添加到任何CNN架构中，提升模型的性能。它总结了之前做过的关于注意力机制应用到图像领域的工作，链接如下：
1.Residual Attention Network for Image Classification
2.SENet：Squeeze-and-Excitation Network
3.SCA-CNN: Spatial and Channel-wise Attention in Convolutional Networks for Image Captioning

其中，受到SENet的启示（不了解SENet的朋友可以看我的文章SENet重点干货及Pytorch实现），CBAM将空间注意力和通道注意力结合起来，实现了两全其美的办法，CNN加入CBAM后能显著提高图像分类和目标检测的准确率。

---

一、什么是CBAM？

如图所示，CBAM是针对feature map进行操作的，通过通道注意力模块和空间注意力模块，使提取到的特征更加精炼，提升模型的表现力。

Channel attention module

通道注意力模块：通道维度不变，压缩空间维度。该模块关注输入图片中有意义的信息(分类任务就关注因为什么分成了不同类别)。

操作流程如下：

输入特征图F，经过两个并行的最大值池化和平均池化将C×H×W的特征图变成C×1×1的大小，然后经过一个Shared MLP(Conv/Linear，ReLU，Conv/Linear)，压缩通道数C/r (reduction=16)，再扩张回C，得到两个激活后的结果。最后将二者相加再接一个sigmoid得到channel_out，channel_out再乘以最初的输入，大小变回了C×H×W，完成通道注意力操作。

此步骤与SENet不同之处是加了一个并行的最大值池化，提取到的高层特征更全面，更丰富。

Spatial attention module

空间注意力模块：空间维度不变，压缩通道维度。该模块关注的是目标的位置信息。

操作流程如下：

将上一步得到的结果通过最大值池化和平均池化分成两个大小为H×W×1的张量，然后通过Concat操作将二者堆叠在一起(C为2)，再通过卷积操作将通道变为1同时保证H和W不变，经过一个sigmoid得到spatial_out，最后spatial_out乘上一步的输入变回C×H×W，完成空间注意力操作

如果看到这还是没有详细理解流程，那就上一张笔者自己画的CBAM详细的流程图供大家参考：

二、Pytorch实现CBAM

CBAM模块完整代码：

import torch
import torch.nn as nn

class CBAMLayer(nn.Module):
    def __init__(self,channel,reduction = 16,spatial_kernel = 7):
        super(CBAMLayer,self).__init__()
        
        #channel attention 压缩H,W为1
        self.max_pool = nn.AdaptiveMaxPool2d(1)
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        
        #shared MLP
        self.mlp = nn.Sequential(
            #Conv2d比Linear方便操作
            #nn.Linear(channel, channel // reduction, bias=False)
            nn.Conv2d(channel,channel//reduction,1,bias = False),
            #inplace=True直接替换，节省内存
            nn.ReLU(inplace=True),
            #nn.Linear(channel // reduction, channel,bias=False)
            nn.Conv2d(channel//reduction,channel,1,bias = False)
        )
        
        #spatial attention
        self.conv = nn.Conv2d(2,1,kernel_size = spatial_kernel,
                              padding = spatial_kernel//2,bias=False)
        self.sigmoid = nn.Sigmoid()
    
    def forward(self,x):
        
        max_out = self.mlp(self.max_pool(x))
        avg_out = self.mlp(self.avg_pool(x))
        channel_out = self.sigmoid(max_out+avg_out)
        x = channel_out * x
        
        max_out,_ = torch.max(x,dim = 1,keepdim=True)
        avg_out = torch.mean(x,dim = 1,keepdim=True)
        spatial_out = self.sigmoid(self.conv(torch.cat([max_out,avg_out],dim = 1)))
        x = spatial_out * x
        return x

测试网络结构代码：

x = torch.randn(1,16,64,64)
net = CBAMLayer(16)
y = net.forward(x)
y.shape

总结

1. CBAM中作者对两个注意力机制使用的先后顺序做了实验，发现通道注意力在空间注意力之前效果最好。
2. 作者在实验中应用了grad-CAM来可视化feature map，是个非常有用的可视化手段，在图像分类任务中可以观察feature map的特征，解释了为什么模型将原图分类到某一类的结果。原文链接：Grad-CAM。该系列方法还有其他改进的方法，感兴趣的朋友可以看一下。
3. 加入CBAM模块不一定会给网络带来性能上的提升，受自身网络还有数据等其他因素影响，甚至会下降。如果网络模型的泛化能力已经很强，而你的数据集不是benchmarks而是自己采集的数据集的话，不建议加入CBAM模块。笔者亲自做实验发现用自己的数据集训练的网络加入CBAM后性能会下降。CBAM性能虽然改进的比SE高了不少，但绝不是无脑加入到网络里就能有提升的。也要根据自己的数据、网络等因素综合考量。总之，就很玄学。。。