[论文理解] CBAM: Convolutional Block Attention Module

Posted 2020-11-13 aoru45

CBAM: Convolutional Block Attention Module

简介

本文利用attention机制，使得针对网络有了更好的特征表示，这种结构通过支路学习到通道间关系的权重和像素间关系的权重，然后乘回到原特征图，使得特征图可以更好的表示。

Convolutional Block Attention Module

这里的结构有点类似与SENet里的支路结构。

对于Channel attention module，先将原feature map分别做global avg pooling 和global max pooling，然后将两pooling后的向量分别连接一个FC层，之后point-wise相加。激活。

这里用global pooling的作用是捕捉全局特征，因为得到的权重描述的是通道间的关系，所以必须要全局特征才能学习到这种关系。

之所以avg pooling和max pooling一起用，是因为作者发现max pooling能够捕捉特征差异，avg pooling能捕捉一般信息，两者一起用的效果要比单独用的实验结果要好，。

结构如图：

对于Spatial attention module，作者使用了1×1的pooling，与上面一样，使用的是1×1的avg pooling和1×1的max pooling，而没有用1×1卷积，两者concat，紧接着是一层7×7卷积，然后激活。最后输出就是1×h×w。

结构如图：

作者提到了两者的顺序，先做channel attention比先做spatial attention要好很多。

后面作者实验了spatial attention module里1×1conv、1×1pooling的效果，最后发现pooing的效果要比卷积的效果要好，因此上面的结构采用的是pooling而不是卷积结构。

后面就是一些结构了。

几句话简单复现了一下。

'''
@Descripttion: This is Aoru Xue's demo,which is only for reference
@version: 
@Author: Aoru Xue
@Date: 2019-09-12 01:24:03
@LastEditors: Aoru Xue
@LastEditTime: 2019-09-12 02:24:25
'''
import torch
import torch.nn as nn

class ChannelAttentionModule(nn.Module):
    def __init__(self,size = 128,r = 2):
        super(ChannelAttentionModule, self).__init__()
        self.max_pooling = nn.MaxPool2d(size)
        self.avg_pooling = nn.AvgPool2d(size)
        self.fc1 = nn.Linear(64,64//r)
        self.fc2 = nn.Linear(64//r,64)
        self.relu = nn.ReLU(inplace=True)
    def forward(self,x):
        max_pool = self.max_pooling(x).view(2,64)
        max_pool = self.fc1(max_pool)
        avg_pool = self.avg_pooling(x).view(2,64)
        avg_pool = self.fc1(avg_pool)
        t = max_pool + avg_pool
        x = self.fc2(t).view(2,64,1,1)
        x = self.relu(x)
        return x
class SpatialAttentionModule(nn.Module):
    def __init__(self,):
        super(SpatialAttentionModule, self).__init__()
        self.conv7x7 = nn.Conv2d(2,64,kernel_size= 7 , stride=1,padding = 3)
        self.sigmoid = nn.Sigmoid()
    def forward(self,x):
        max_pool = torch.max(x,dim = 1)[0]
        avg_pool = torch.mean(x,dim = 1)
        x = self.conv7x7(torch.stack([max_pool,avg_pool],dim = 1))
        x = self.sigmoid(x)
        return x
class ResBlock(nn.Module):
    def __init__(self,):
        super(ResBlock, self).__init__()
        self.channel_module = ChannelAttentionModule(r = 2)
        self.spatial_module = SpatialAttentionModule()
    def forward(self,x):
        c = self.channel_module(x)
        x = c*x
        s = self.spatial_module(x)
        x = s * x
        return x
if __name__ == "__main__":
    x = torch.randn(2,64,128,128)
    net = ResBlock()
    print(net(x).size())