概述

本文提出了一个注意力模块，融合空间注意力和通道注意力，因为是注意力机制，所以不会对数据的shape进行变化，只是从通道维度和空间维度，对其中重要的内容进行凸显，对其中不重要的内容进行抑制，所以可以用于网络的各层。

细节

这张是CBAM模块的总体概述，先对特征图做通道注意力，再做空间注意力

下面是通道注意力的细节图，我们可以看到这里的实现其实是和之前的SE block差不多的，SE block中只使用了全局平均池化，这里多用了一个全局最大池化。

空间注意力也是，相对于之前的方式多了一个取最大值的操作。

实现

paddle版

import paddle
import paddle.nn as nn

# 通道注意力机制
class ChannelAttentionModule(nn.Layer):
    def __init__(self, in_channels, reduction=16):
        super(ChannelAttentionModule, self).__init__()
        # 全局平均池化和全局最大池化 将尺寸变为1*1 但是通道数不变
        self.avg_pool = nn.AdaptiveAvgPool2D(1)
        self.max_pool = nn.AdaptiveMaxPool2D(1)
        self.shared_MLP = nn.Sequential(
            nn.Linear(in_channels, in_channels // reduction),
            nn.ReLU(),
            nn.Linear(in_channels // reduction, in_channels),
        )
        self.sigmoid=nn.Sigmoid()
    def forward(self,x):
        # x:[n,c,h,w]
        # 我们不需要w,h 因为池化操作之后就是 1*1了
        n, c, _, _ = x.shape
        avg_out = self.avg_pool(x).flatten(1)  # avg_out:[n,c*h*w]=>[n,c]
        max_out = self.max_pool(x).flatten(1)  # max_out:[n,c*h*w]=>[n,c]
        avg_out = self.shared_MLP(avg_out).reshape([n, c, 1, 1])  # avg_out:[n,c,1,1]
        max_out = self.shared_MLP(max_out).reshape([n, c, 1, 1])  # max_out:[n,c,1,1]

        return self.sigmoid(avg_out+max_out)

# 空间注意力模块
class SpatialAttentionModule(nn.Layer):
    def __init__(self):
        super(SpatialAttentionModule, self).__init__()
        self.conv2d = nn.Conv2D(in_channels=2, out_channels=1, kernel_size=7, stride=1, padding=3)
        self.sigmoid = nn.Sigmoid()

    def forward(self,x):
        # x:[n,c,h,w]
        # mean和max会沿着通道进行求平均和求最大值，并且我们保留了通道这个维度，不保留的话就是[n,h,w]了
        avg_out=paddle.mean(x,axis=1,keepdim=True) # avg_out:[n,1,h,w]
        max_out=paddle.max(x,axis=1,keepdim=True) # max_out:[n,1,h,w]
        out=paddle.concat([avg_out,max_out],axis=1) # out:[n,2,h,w]
        out=self.conv2d(out)
        out=self.sigmoid(out)
        return out

class CBDM(nn.Layer):
    def __init__(self, in_channels, reduction=16):
        super(CBDM, self).__init__()
        self.channel_attention = ChannelAttentionModule(in_channels, reduction)
        self.spatial_attention = SpatialAttentionModule()

    def forward(self,x):
        out=self.channel_attention(x)*x
        out=self.spatial_attention(out)*x
        return out


def main():
    x=paddle.randn([3,16,15,15])
    cbdm=CBDM(16)
    print(cbdm(x).shape)

if __name__ == '__main__':
    main()

torch版

import torch
import torch.nn as nn

# 通道注意力机制
class ChannelAttentionModule(nn.Module):
    def __init__(self, in_channels, reduction=16):
        super(ChannelAttentionModule, self).__init__()
        # 全局平均池化和全局最大池化 将尺寸变为1*1 但是通道数不变
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.max_pool = nn.AdaptiveAvgPool2d(1)
        self.shared_MLP = nn.Sequential(
            nn.Linear(in_channels, in_channels // reduction),
            nn.ReLU(),
            nn.Linear(in_channels // reduction, in_channels),
        )
        self.sigmoid=nn.Sigmoid()
    def forward(self,x):
        # x:[n,c,h,w]
        # 我们不需要w,h 因为池化操作之后就是 1*1了
        n, c, _, _ = x.shape
        avg_out = self.avg_pool(x).flatten(1)  # avg_out:[n,c*h*w]=>[n,c]
        max_out = self.max_pool(x).flatten(1)  # max_out:[n,c*h*w]=>[n,c]
        avg_out = self.shared_MLP(avg_out).reshape([n, c, 1, 1])  # avg_out:[n,c,1,1]
        max_out = self.shared_MLP(max_out).reshape([n, c, 1, 1])  # max_out:[n,c,1,1]

        return self.sigmoid(avg_out+max_out)

# 空间注意力模块
class SpatialAttentionModule(nn.Module):
    def __init__(self):
        super(SpatialAttentionModule, self).__init__()
        self.conv2d = nn.Conv2d(in_channels=2, out_channels=1, kernel_size=7, stride=1, padding=3)
        self.sigmoid = nn.Sigmoid()

    def forward(self,x):
        # x:[n,c,h,w]
        # mean和max会沿着通道进行求平均和求最大值，并且我们保留了通道这个维度，不保留的话就是[n,h,w]了
        avg_out=torch.mean(x,dim=1,keepdim=True) # avg_out:[n,1,h,w]
        max_out=torch.max(x,dim=1,keepdim=True)[0] # max_out:[n,1,h,w]
        out=torch.cat((avg_out,max_out),1) # out:[n,2,h,w]
        out=self.conv2d(out)
        out=self.sigmoid(out)
        return out

class CBDM(nn.Module):
    def __init__(self, in_channels, reduction=16):
        super(CBDM, self).__init__()
        self.channel_attention = ChannelAttentionModule(in_channels, reduction)
        self.spatial_attention = SpatialAttentionModule()

    def forward(self,x):
        out=self.channel_attention(x)*x
        out=self.spatial_attention(out)*x
        return out


def main():
    x=torch.randn([3,16,15,15])
    cbdm=CBDM(16)
    print(cbdm(x).shape)

if __name__ == '__main__':
    main()