快捷搜索:

【注意力机制集锦】Channel Attention通道注意力网络结构、源码解读系列一

Channel Attention网络结构、源码解读系列一

SE-Net、SK-Net与CBAM

1 SENet

1.1 Squeeze-and-Excitation Blocks

SE Block模块主要由Squeeze操作和Excitation操作组成:Squeeze操作负责将spatial维度进行全局池化(比如7 x 7 -->1 x 1);Excitation操作则学习池化后的通道依赖关系,并进行通道权重的赋权。上图网络结构其实很好地概括了SENet的主题思想,下面我将会从Squeeze和Excitation两个方面具体讲解。

1.1.1 Squeeze: Global Information Embedding

网络结构最开始的部分Ftr:X->U是以往的经典卷积结构,U之后的部分才是SENet的创新部分:使用全局平均池化在H和W两个维度对U进行Squeeze,将一个channel上整个空间特征编码为一个全局特征,得到1x1xC的中间输出。说得通俗点,这里其实就是使用一个二维的池化核对特征图进行降维,由原来的H、W、C上的3个维度降到了C这1个维度上,使得后续的通道赋权操作可行,其公式如下图所示:

1.1.2 Excitation: Adaptive Recalibration

1.1.3 举个栗子:SE-ResNet Module

1.2 代码实现

1.2.1 SE module

SE的实现如下代码所示,具体每一步我都做了详细的注释。如果前面的公式看不明白,对应这里的函数操作可能会帮助理解公式。

class SELayer(nn.Module):
    def __init__(self, channel, reduction=16):
        super(SELayer, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)# Squeeze操作的定义
        self.fc = nn.Sequential(# Excitation操作的定义
            nn.Linear(channel, channel // reduction, bias=False),# 压缩
            nn.ReLU(inplace=True),
            nn.Linear(channel // reduction, channel, bias=False),# 恢复
            nn.Sigmoid()# 定义归一化操作
        )

    def forward(self, x):
        b, c, _, _ = x.size()# 得到H和W的维度,在这两个维度上进行全局池化
        y = self.avg_pool(x).view(b, c)# Squeeze操作的实现
        y = self.fc(y).view(b, c, 1, 1)# Excitation操作的实现
        # 将y扩展到x相同大小的维度后进行赋权
        return x * y.expand_as(x)

1.2.2 SE-ResNet

class SEBasicBlock(nn.Module):
    expansion = 1

    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
                 base_width=64, dilation=1, norm_layer=None,
                 *, reduction=16):
        super(SEBasicBlock, self).__init__()
        self.conv1 = conv3x3(inplanes, planes, stride)
        self.bn1 = nn.BatchNorm2d(planes)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = conv3x3(planes, planes, 1)
        self.bn2 = nn.BatchNorm2d(planes)
        self.se = SELayer(planes, reduction)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        residual = x
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)
        out = self.se(out)# 加入通道注意力机制

        if self.downsample is not None:
            residual = self.downsample(x)

        out += residual
        out = self.relu(out)

        return out

2 SKNet

CVPR2019的文章Selective Kernel Networks,这篇文章也是致敬了SENet的思想。 SENet提出了Sequeeze and Excitation block,而SKNet提出了Selective Kernel Convolution. 二者都可以很方便的嵌入到现在的网络结构,比如ResNet、Inception、ShuffleNet,实现精度的提升。

2.1 Selective Kernel Convolution

2.1.1 Split

对输入X使用不同的卷积核生成不同的特征输出,上图所示的是使用3x3和5x5的卷积核进行的卷积操作,为了提高运算效率,5x5的卷积操作是用空洞率为2、卷积核为3x3的空洞卷积实现的,并且使用了分组卷积、深度可分离卷积、BatchNorm和ReLU。

2.1.2 Fuze

将得到的多个特征输出进行信息融合,即pytorch中的sum操作,得到新的特征图U,即下图中的公式(1);然后利用Squeeze相同的操作生成通道这个维度的信息,即下图中的公式(2);最后利用1层全连接层FC学习通道之间的依赖关系,最后使用ReLU和BatchNorm进行归一,即下图中的公式(3)。相关公式如下:

2.1.3 Select

在通道这个维度对多个分支得到的最终特征图进行赋权,使用sigmoid函数。最后将所有分支加权后的特征图相加,得到最终的输出。

2.2 代码实现

结合上述的讲解,代码其实很明了,具体的定义及操作我都作了注释,可以参看注释进行理解。

class SKConv(nn.Module):
    def __init__(self, features, WH, M, G, r, stride=1 ,L=32):
        super(SKConv, self).__init__()
        d = max(int(features/r), L)
        self.M = M
        self.features = features
        self.convs = nn.ModuleList([])
        # 生成M个分支,将其添加到convs中,每个分支采用不同的卷积核和不同规模的padding,保证最终得到的特征图大小一致
        for i in range(M):
            self.convs.append(nn.Sequential(
                nn.Conv2d(features, features, kernel_size=3+i*2, stride=stride, padding=1+i, groups=G),
                nn.BatchNorm2d(features),
                nn.ReLU(inplace=False)
            ))
        # 学习通道间依赖的全连接层
        self.fc = nn.Linear(features, d)
        self.fcs = nn.ModuleList([])
        for i in range(M):
            self.fcs.append(
                nn.Linear(d, features)
            )
        self.softmax = nn.Softmax(dim=1)
        
    def forward(self, x):
        for i, conv in enumerate(self.convs):
            fea = conv(x).unsqueeze_(dim=1)
            if i == 0:
                feas = fea
            else:
                feas = torch.cat([feas, fea], dim=1)
        fea_U = torch.sum(feas, dim=1)# 将多个分支得到的特征图进行融合
        fea_s = fea_U.mean(-1).mean(-1)# 在channel这个维度进行特征抽取
        fea_z = self.fc(fea_s)# 学习通道间的依赖关系
        # 赋权操作,由于是对多维数组赋权,所以看起来比SENet麻烦一些
        for i, fc in enumerate(self.fcs):
            vector = fc(fea_z).unsqueeze_(dim=1)
            if i == 0:
                attention_vectors = vector
            else:
                attention_vectors = torch.cat([attention_vectors, vector], dim=1)
        attention_vectors = self.softmax(attention_vectors)
        attention_vectors = attention_vectors.unsqueeze(-1).unsqueeze(-1)
        fea_v = (feas * attention_vectors).sum(dim=1)
        return fea_v

3 CBAM

CBAM( Convolutional Block Attention Module )是一种轻量化的通道注意力机制,也是目前应用比较广泛的一种视觉注意力机制,在2018年的ECCV中提出。文章同时使用了Channel Attention和Spatial Attention,发现将两种attention串联在一起效果较好。

3.1 Convolutional Block Attention Module

下图是CBAM的网络结构图。

可以看到CBAM包含2个独立的子模块, 通道注意力模块(Channel Attention Module,CAM) 和空间注意力模块(Spartial Attention Module,SAM) ,分别进行通道与空间上的赋权。这样不只能够节约参数和计算力,并且保证了其能够做为即插即用的模块集成到现有的网络架构中去。

3.1.1 Channel attention module

通道注意力机制的基本思想与SENet相同,但是具体操作与SENet略有不同,不同部分我用红色进行了标记。 首先,将输入的特征图F(H×W×C)分别经过基于H和W两个维度的 全局最大池化(MaxPool)和全局平均池化(AvgPool),得到两个1×1×C的特征图; 然后,将两个特征图送入一个 共享权值的双层神经网络(MLP)进行通道间依赖关系的学习,两层神经层之间通过压缩比r实现降维。 最后,将MLP输出的特征进行基于element-wise的加和操作,再经过sigmoid激活操作,生成最终的通道加权,即M_c。其公式如下图:

3.1.2 Spatial attention module

3.2 代码实现

3.2.1 CA&SA

具体的网络定义及操作实现参见我的代码注释。

class ChannelAttention(nn.Module):
    def __init__(self, in_planes, ratio=16):
        super(ChannelAttention, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)# 定义全局平均池化
        self.max_pool = nn.AdaptiveMaxPool2d(1)# 定义全局最大池化
        # 定义CBAM中的通道依赖关系学习层,注意这里是使用1x1的卷积实现的,而不是全连接层
        self.fc = nn.Sequential(nn.Conv2d(in_planes, in_planes // 16, 1, bias=False),
                               nn.ReLU(),
                               nn.Conv2d(in_planes // 16, in_planes, 1, bias=False))
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        avg_out = self.fc(self.avg_pool(x))# 实现全局平均池化
        max_out = self.fc(self.max_pool(x))# 实现全局最大池化
        out = avg_out + max_out# 两种信息融合
        # 最后利用sigmoid进行赋权
        return self.sigmoid(out)

class SpatialAttention(nn.Module):
    def __init__(self, kernel_size=7):
        super(SpatialAttention, self).__init__()
        # 定义7*7的空间依赖关系学习层
        self.conv1 = nn.Conv2d(2, 1, kernel_size, padding=kernel_size//2, bias=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        avg_out = torch.mean(x, dim=1, keepdim=True)# 实现channel维度的平均池化
        max_out, _ = torch.max(x, dim=1, keepdim=True)# 实现channel维度的最大池化
        x = torch.cat([avg_out, max_out], dim=1)# 拼接上述两种操作的到的两个特征图
        x = self.conv1(x)# 学习空间上的依赖关系
        # 对空间元素进行赋权
        return self.sigmoid(x)

3.2.2 CBAM_ResNet

篇幅限制,这里仅展示BasicBlock的CA&SA的添加

class BasicBlock(nn.Module):
    expansion = 1

    def __init__(self, inplanes, planes, stride=1, downsample=None):
        super(BasicBlock, self).__init__()
        self.conv1 = conv3x3(inplanes, planes, stride)
        self.bn1 = nn.BatchNorm2d(planes)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = conv3x3(planes, planes)
        self.bn2 = nn.BatchNorm2d(planes)
		# 定义ca和sa,注意CA与channel num有关,需要指定这个超参!!!
        self.ca = ChannelAttention(planes)
        self.sa = SpatialAttention()
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        residual = x
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)
        out = self.conv2(out)
        out = self.bn2(out)
        out = self.ca(out) * out# 对channel赋权
        out = self.sa(out) * out# 对spatial赋权
        if self.downsample is not None:
            residual = self.downsample(x)
        out += residual
        out = self.relu(out)

        return out
经验分享 程序员 微信小程序 职场和发展