【视觉注意力机制】SE、CBAM、ECA三种可插拔注意力模块结构实现与详解

SE、CBAM 以及 ECA 三种注意力机制的结构实现与代码详解如下所示。

代码可参考：https://github.com/XuecWu/External-Attention-pytorch

import torch
import torch.nn as nn
import math

#----------------------------#
# SE注意力机制
#----------------------------#
class se_block(nn.Module):
    def __init__(self, channel, ratio=16):
        super(se_block, self).__init__()
        #--------------------------------------------------#
        # 此为自适应的二维平均全局池化操作
        # 通道数不会发生改变
        # The output is of size H x W, for any input size.
        # AdaptiveAvgPool2d(1) = AdaptiveAvgPool2d((1,1))
        #--------------------------------------------------#
        self.avg_pool = nn.AdaptiveAvgPool2d(1)

        #-------------------------------------------------------#
        # 对于全连接层,第一个参数为输入的通道数,第二个参数为输入的通道数
        # 之后经过ReLU来提升模型的非线性表达能力,以及对特征信息进行编码
        # sigmoid还是一个激活函数,来提升模型的非线性表达能力
        # ratio越大其对于特征融合以及信息表达所产生的影响越大,
        # 压缩降维对于学习通道之间的依赖关系有着不利影响
        #-------------------------------------------------------#
        self.fc = nn.Sequential(
            nn.Linear(channel, channel // ratio, bias=False),
            nn.ReLU(inplace=True),
            nn.Linear(channel // ratio, channel, bias=False),
            nn.Sigmoid()
        )

    def forward(self, x):

        b, c, _, _ = x.size()
        #-----------------------------------------------------------------#
        # 先进行自适应二维全局平均池化,然后进行一个reshape的操作
        # 之后使其通过一个全连接层、一个ReLU、一个全连接层、一个Sigmoid层
        # 再将其reshape成之前的shape即可
        # 最后将注意力权重y和输入X按照通道加权相乘，调整模型对输入x不同通道的重视程度
        #------------------------------------------------------------=----#
        y = self.avg_pool(x).view(b, c)
        y = self.fc(y).view(b, c, 1, 1)
        return x * y


#--------------------------------------#
# CBAM注意力机制 包含通道注意力以及空间注意力
#--------------------------------------#
class ChannelAttention(nn.Module):
    def __init__(self, in_planes, ratio=8):
        super(ChannelAttention, self).__init__()

        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.max_pool = nn.AdaptiveMaxPool2d(1)
        #-----------------------------------------#
        # 利用1x1卷积代替全连接,以减小计算量以及模型参数
        # 整体前向传播过程为 卷积 ReLU 卷积 进而实现特征
        # 信息编码以及增强网络的非线性表达能力
        #-----------------------------------------#
        self.fc1   = nn.Conv2d(in_planes, in_planes // ratio, 1, bias=False)
        self.relu1 = nn.ReLU()
        self.fc2   = nn.Conv2d(in_planes // ratio, in_planes, 1, bias=False)
        #----------------------------------------#
        # 定义sigmoid激活函数以增强网络的非线性表达能力
        #----------------------------------------#
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        #-------------------------------------#
        # 分成两部分,一部分为平均池化,一部分为最大池化
        # 之后将两部分的结果相加再经过sigmoid作用
        #-------------------------------------#
        avg_out = self.fc2(self.relu1(self.fc1(self.avg_pool(x))))
        max_out = self.fc2(self.relu1(self.fc1(self.max_pool(x))))
        out     = avg_out + max_out
        return self.sigmoid(out)


class SpatialAttention(nn.Module):
    def __init__(self, kernel_size=7):
        super(SpatialAttention, self).__init__()

        assert kernel_size in (3, 7), 'kernel size must be 3 or 7'
        padding      = 3 if kernel_size == 7 else 1
        #-------------------------------------#
        # 这个卷积操作为大核卷积操作,其虽然可以计算
        # 空间注意力但是仍无法有效建模远距离依赖关系
        #-------------------------------------#
        self.conv1   = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        #----------------------------------------------------#
        # 整体前向传播过程即为先分别做平均池化操作,再做最大池化操作
        # 其中的dim：
        # 指定为1时,求得是列的平均值
        # 指定为0时,求得是行的平均值
        # 之后将两个输出按照列维度进行拼接,此时通道数为2
        # 拼接之后通过一个大核卷积将双层特征图转为单层特征图,此时通道为1
        # 最后通过sigmoid来增强模型的非线性表达能力
        #-----------------------------------------------------#
        avg_out    = torch.mean(x, dim=1, keepdim=True)
        max_out, _ = torch.max(x, dim=1, keepdim=True)
        x          = torch.cat([avg_out, max_out], dim=1)
        x          = self.conv1(x)
        return self.sigmoid(x)


class cbam_block(nn.Module):
    def __init__(self, channel, ratio=8, kernel_size=7):
        super(cbam_block, self).__init__()
        #----------------------------#
        # 定义好通道注意力以及空间注意力
        #----------------------------#
        self.channelattention = ChannelAttention(channel, ratio=ratio)
        self.spatialattention = SpatialAttention(kernel_size=kernel_size)

        #----------------------------#
        # 输入x先与通道注意力权重相乘
        # 之后将输出与空间注意力权重相乘
        #----------------------------#
    def forward(self, x):
        x = x * self.channelattention(x)
        x = x * self.spatialattention(x)
        return x


#----------------------------#
# ECA注意力机制
#----------------------------#
class eca_block(nn.Module):
    def __init__(self, channel, b=1, gamma=2):
        super(eca_block, self).__init__()
        #----------------------------------#
        # 根据通道数求出卷积核的大小kernel_size
        #----------------------------------#
        kernel_size = int(abs((math.log(channel, 2) + b) / gamma))
        kernel_size = kernel_size if kernel_size % 2 else kernel_size + 1

        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.conv     = nn.Conv1d(1, 1, kernel_size=kernel_size, padding=(kernel_size - 1) // 2, bias=False)
        self.sigmoid  = nn.Sigmoid()

    def forward(self, x):
        #------------------------------------------#
        # 显示全局平均池化,再是k*k的卷积,
        # 最后为Sigmoid激活函数,进而得到每个通道的权重w
        # 最后进行回承操作,得出最终结果
        #------------------------------------------#
        y = self.avg_pool(x)
        y = self.conv(y.squeeze(-1).transpose(-1, -2)).transpose(-1, -2).unsqueeze(-1)
        y = self.sigmoid(y)
        return x * y.expand_as(x)