【深度学习注意力机制系列】—— SKNet注意力机制（附pytorch实现）

news2024/9/30 1:34:18

SKNet（Selective Kernel Network）是一种用于图像分类和目标检测任务的深度神经网络架构，其核心创新是引入了选择性的多尺度卷积核（Selective Kernel）以及一种新颖的注意力机制，从而在不增加网络复杂性的情况下提升了特征提取的能力。SKNet的设计旨在解决多尺度信息融合的问题，使网络能够适应不同尺度的特征。

1. 核心思想

SKNet的核心思想是**通过选择性地应用不同尺度的卷积核，从而在不同层级上捕捉多尺度特征。**为了实现这一点，SKNet引入了一个选择模块，用于自适应地决定在每个通道上使用哪些尺度的卷积核。这种选择性的多尺度卷积核有助于提升特征表示的能力，使网络更具适应性和泛化能力。

2. 结构

SKNet的结构如下：

在这里插入图片描述

实现机制：

split：对特征图进行多分支分离卷积，各分支使用不同的卷积核（感受野不同）进行特征提取。（并未对原始特征图进行拆解分离，只是使用不同的卷积核对原始特征图进行卷积操作）。假设分支为n,则特征图维度变换为 (c, h, w) -> (n, c, h, w)，原文中n=2。
Fuse：将多个分支的特征图提取结果相加。特征图维度变换为 (n, c, h, w) -> (c, h, w)。再通过全局平均池，特征图维度变换为 (c, h, w) -> (c, 1, 1)，然后利用全连接层进行降维（限制了最低维度，通过全连接层生成d×1的向量(图中的z)，公式如图中所示(δ表示ReLU激活函数，B表示Batch Noramlization,W是一个d×C的维的)。d的取值是由公式d = max(C/r,L)确定，r是一个缩小的比率(与SENet中相似)，L表示d的最小值，原文实验中L的值为32。），再利用两个（或多个，和分支数目相同，原论文中为两个）全连接层进行升维，得到两个（多个）维度同降维前相同的特征图（向量）。在对两个特征向量进行softmax处理。假设分支为n,则特征图维度为 n个(c, 1, 1) ，原文中n=2，即a->(c, 1, 1)， b->(c, 1, 1)。
select：利用softmax处理后的多个特征向量分别乘以第一步中的多分支提取的特征图结果。特征维度变化为n个(c, 1 ,1) * n 个(c, h ,w) = (n, c, h, w)。最后将n个特征图进行相加。

3. 优势

SKNet的设计在以下几个方面具有优势：

多尺度信息融合

通过选择性地应用不同尺度的卷积核，SKNet能够有效地融合多尺度的特征信息。这有助于网络捕捉不同层次的视觉特征，提高了特征的表征能力。

自适应性

选择模块使网络能够自适应地选择卷积核的尺度，从而适应不同任务和图像的特点。这种自适应性能够使网络在各种场景下都能表现出色。

减少计算成本

尽管引入了多尺度卷积核，但由于选择模块的存在，SKNet只会选择一部分卷积核进行计算，从而减少了计算成本，保持了网络的高效性。

4.代码实现

class SKNet(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(SKNet, self).__init__()
        d = max(in_channels // r, L)  
        self.M = M
        self.out_channels = out_channels
        self.conv = nn.ModuleList() 
        for i in range(M):
            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) 
        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) 
    def forward(self, input):
        batch_size = input.size(0)
        output = []
        for i, conv in enumerate(self.conv):
            output.append(conv(input))
        U = reduce(lambda x, y: x + y, output)  
        s = self.global_pool(U)  
        z = self.fc1(s)
        a_b = self.fc2(z) 
        a_b = a_b.reshape(batch_size, self.M, self.out_channels, -1) 
        a_b = self.softmax(a_b) 
        a_b = list(a_b.chunk(self.M, dim=1))  
        a_b = list(map(lambda x: x.reshape(batch_size, self.out_channels, 1, 1),
                       a_b))  
        V = list(map(lambda x, y: x * y, output,
                     a_b))  
        V = reduce(lambda x, y: x + y,
                   V)  
        return V