SK 注意力模块 原理分析与代码实现

Posted 2023-01-06 一颗小树x

前言

本文介绍SK模块，一种通道注意力模块，它是在SK-Nets中提出的，SK-Nets是2019 CVPR中的论文；SK模块可以被用于CV模型中，能提取模型精度，所以给大家介绍一下它的原理，设计思路，代码实现，如何应用在模型中。

一、SK 注意力模块

SK 注意力模块，它是一种通道注意力模块；常常被应用与视觉模型中。支持即插即用，即：它能对输入特征图进行通道特征加强，而且最终SK模块输出，不改变输入特征图的大小。

• 背景：它在SE注意力模型的基础上进行改进；SK在学习通道注意力信息中，通过网络自己学习，来选择融合不同感受野的特征图信息；
• 设计：SK是Selective Kernel的缩写，意思是选择卷积核，它能使用不同尺寸的卷积核，对输入特征图进行提取特征，然后自动选择合适的卷积核所提取出来的信息。
• 作用：不同大小的感受视野（卷积核），对于不同尺度（远近、大小）的目标会有不同的效果。
• 思想：对输入特征图中的不同大小物体，自动选择对应的卷积核来提取特征，生成通道注意力信息。

 下面分析一下，SK是如何实现通道注意力的；首先看一下，模块的结构：

  Sk模块的流程思路如下:

1、首先输入特征图，它的维度是C*H*W；分别使用3x3、5x5的卷积核，对特征图进行卷积操作；

         细节1：对于kernel = 3 * 3，使用dilation = 1，得到的特征依然是C*H*W；

        细节2：对于kernel = 5 * 5，使用dilation = 2，得到的特征也是C*H*W；（为了减轻计算，还可以使用2个 3 * 3卷积核，代替1个5 * 5卷积核，来对特征图进行卷积操作）

2、得到相同尺寸的特征图，进行融合，操作是：对特征图把相应像素相加；得到的特征图U，维度是C*H*W；

3、对融合后的特征图，进行空间维度的全局平均池化操作，得到的特征图S，维度是C*1*1；

4、通过FC全连接提取通道注意力信息；得到的特征图Z，维度是d*1*1；d小于C。

        细节1：这里在做通道注意力信息提取时，是一个降维操作，及d<C；后里又把通道数由d变换为C。

5、前面得到的通道注意力信息Z，分别与卷积核A、卷积核B，进行卷积操作；再经过Softmax处理；得到了每个卷积核对应的通道注意力信息；

        细节1：分别使用不同的卷积核，与通道注意力信息Z进行卷积；上面图有2个卷积核（A、B）；

        细节2：为什么要经过Softmax处理呢？理论上，不同的卷积核与通道注意力信息Z进行卷积，也能得到了每个卷积核对应的通道注意力信息；

        Softmax的作用是使得通道注意力信息a、通道注意力信息b，两者的权重相加合计为1。即，在多个卷积核中，选择权重较大的那个卷积核，认为它比较合适。

6、得到通道注意力信息a (C*1*1)，与前面经过卷积核A处理的特征图(C*H*W)，进行逐通道乘，最终输出具有卷积核A——通道注意力的特征图。(C*H*W)

        得到通道注意力信息b (C*1*1)，与前面经过卷积核B处理的特征图(C*H*W)，进行逐通道乘，最终输出具B——通道注意力的特征图。(C*H*W)

7、最后，将卷积核A、卷积核B通道注意力的特征图，进行融合（两幅特征图，对应像素相加）；得到最终的通道注意力特征图V(C*H*W)

下图是SK模块的计算的对应公式：

​

参考：https://blog.csdn.net/frighting_ing/article/details/121429665

二、代码实现

ECA 通道注意力模块，基于pytorch版本的代码如下：

import torch.nn as nn
import torch
from functools import reduce
class SKConv(nn.Module):
    def __init__(self,in_channels,out_channels,stride=1,M=2,r=16,L=32):
        '''
        :param in_channels:  输入通道维度
        :param out_channels: 输出通道维度   原论文中 输入输出通道维度相同
        :param stride:  步长，默认为1
        :param M:  分支数
        :param r: 特征Z的长度，计算其维度d 时所需的比率（论文中 特征S->Z 是降维，故需要规定 降维的下界）
        :param L:  论文中规定特征Z的下界，默认为32
        采用分组卷积： groups = 32,所以输入channel的数值必须是group的整数倍
        '''
        super(SKConv,self).__init__()
        d=max(in_channels//r,L)   # 计算从向量C降维到 向量Z 的长度d
        self.M=M
        self.out_channels=out_channels
        self.conv=nn.ModuleList()  # 根据分支数量 添加 不同核的卷积操作
        for i in range(M):
            # 为提高效率，原论文中扩张卷积5x5为（3X3，dilation=2）来代替。
            self.conv.append(nn.Sequential(nn.Conv2d(in_channels,out_channels,3,stride,padding=1+i,dilation=1+i,groups=32,bias=False),
                                           nn.BatchNorm2d(out_channels),
                                           nn.ReLU(inplace=True)))
        self.global_pool=nn.AdaptiveAvgPool2d(output_size = 1) # 自适应pool到指定维度    这里指定为1，实现 GAP
        self.fc1=nn.Sequential(nn.Conv2d(out_channels,d,1,bias=False),
                               nn.BatchNorm2d(d),
                               nn.ReLU(inplace=True))   # 降维
        self.fc2=nn.Conv2d(d,out_channels*M,1,1,bias=False)  # 升维
        self.softmax=nn.Softmax(dim=1) # 指定dim=1  使得两个全连接层对应位置进行softmax,保证 对应位置a+b+..=1
    def forward(self, input):
        batch_size=input.size(0)
        output=[]
        #the part of split
        for i,conv in enumerate(self.conv):
            #print(i,conv(input).size())
            output.append(conv(input))    #[batch_size,out_channels,H,W]
        #the part of fusion
        U=reduce(lambda x,y:x+y,output) # 逐元素相加生成 混合特征U  [batch_size,channel,H,W]
        print(U.size())            
        s=self.global_pool(U)     # [batch_size,channel,1,1]
        print(s.size())
        z=self.fc1(s)  # S->Z降维   # [batch_size,d,1,1]
        print(z.size())
        a_b=self.fc2(z) # Z->a，b 升维  论文使用conv 1x1表示全连接。结果中前一半通道值为a,后一半为b   [batch_size,out_channels*M,1,1]
        print(a_b.size())
        a_b=a_b.reshape(batch_size,self.M,self.out_channels,-1) #调整形状，变为 两个全连接层的值[batch_size,M,out_channels,1]  
        print(a_b.size())
        a_b=self.softmax(a_b) # 使得两个全连接层对应位置进行softmax [batch_size,M,out_channels,1]  
        #the part of selection
        a_b=list(a_b.chunk(self.M,dim=1))#split to a and b   chunk为pytorch方法，将tensor按照指定维度切分成 几个tensor块 [[batch_size,1,out_channels,1],[batch_size,1,out_channels,1]
        print(a_b[0].size())
        print(a_b[1].size())
        a_b=list(map(lambda x:x.reshape(batch_size,self.out_channels,1,1),a_b)) # 将所有分块  调整形状，即扩展两维  [[batch_size,out_channels,1,1],[batch_size,out_channels,1,1]
        V=list(map(lambda x,y:x*y,output,a_b)) # 权重与对应  不同卷积核输出的U 逐元素相乘[batch_size,out_channels,H,W] * [batch_size,out_channels,1,1] = [batch_size,out_channels,H,W]
        V=reduce(lambda x,y:x+y,V) # 两个加权后的特征 逐元素相加  [batch_size,out_channels,H,W] + [batch_size,out_channels,H,W] = [batch_size,out_channels,H,W]
        return V    # [batch_size,out_channels,H,W]

三、SK应用在模型中

sK模块可以被用于CV模型中，能较有效提取模型精度；它是即插即用的，用法和SE模型差不多的。

应用示例1：

在主干网络（Backbone）中，加入Sk模块，加强通道特征，提高模型性能；

应用示例2：

在主干网络（Backbone）末尾，加入Sk模型，加强整体的通道特征，提高模型性能；

应用实例3：

在多尺度特征分支中，加入Sk模块，加强加强通道特征，提高模型性能。

相关参考信息

论文名称：Selective Kernel Networks

论文链接：https://arxiv.org/abs/1903.06586

论文代码：https://github.com/hujie-frank/SENet

本文只供大家参考与学习，谢谢~

后面还会介绍其它注意力模型：CBAM、DANet、CA等注意力模块。

相关文章推荐：

视觉 注意力机制——通道注意力、空间注意力、自注意力_一颗小树x的博客-CSDN博客_空间注意力模块和通道注意力模块

SE 注意力模块 原理分析与代码实现_一颗小树x的博客-CSDN博客_se注意力模块

ECA 注意力模块 原理分析与代码实现_一颗小树x的博客-CSDN博客_eca模块