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

时间:2022-10-15 19:56:39

前言

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

一、SK 注意力模块

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

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

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

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模块的计算的对应公式:

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模块