pytorch代码实现注意力机制之SK Attention

最新推荐文章于 2025-03-22 21:23:00 发布

原创最新推荐文章于 2025-03-22 21:23:00 发布 · 1k 阅读

12 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #人工智能 #python #目标检测 #深度学习 #视觉检测 #计算机视觉

pytorch实现各种注意力机制专栏收录该内容

26 篇文章

订阅专栏

本文介绍了一种在卷积神经网络中引入的SK注意力机制，通过动态调整神经元的感受野大小，以适应不同尺度的输入信息。实验结果表明，SKNet在保持较低模型复杂性的同时，优于现有最先进的架构在ImageNet和CIFAR数据集上表现优秀。

SK注意力机制

简介：在标准卷积神经网络（CNN）中，每一层人工神经元的感受野被设计为共享相同的大小。在神经科学界众所周知，视觉皮层神经元的感受野大小受到刺激的调节，这在构建 CNN 时很少被考虑。我们在CNN中提出了一种动态选择机制，该机制允许每个神经元根据输入信息的多个尺度自适应地调整其感受野大小。设计了一个称为选择性内核（SK）单元的构建块，其中使用由这些分支中的信息引导的 softmax 注意力融合具有不同内核大小的多个分支。对这些分支的不同关注会产生不同大小的融合层神经元的有效感受野。多个 SK 单元堆叠到称为选择性内核网络（SKNets）的深度网络中。在ImageNet和CIFAR基准测试中，我们通过经验表明，SKNet以较低的模型复杂性优于现有的最先进的架构。详细分析表明，SKNet中的神经元可以捕获不同尺度的目标对象，这验证了神经元根据输入自适应调整其感受野大小的能力。

结构图

代码实现

import numpy as np
import torch
from torch import nn
from torch.nn import init
from collections import OrderedDict


class SKAttention(nn.Module):

    def __init__(self, channel=512, kernels=[1, 3, 5, 7], reduction=16, group=1, L=32):
        super().__init__()
        self.d = max(L, channel // reduction)
        self.convs = nn.ModuleList([])
        for k in kernels:
            self.convs.append(
                nn.Sequential(OrderedDict([
                    ('conv', nn.Conv2d(channel, channel, kernel_size=k, padding=k // 2, groups=group)),
                    ('bn', nn.BatchNorm2d(channel)),
                    ('relu', nn.ReLU())
                ]))
            )
        self.fc = nn.Linear(channel, self.d)
        self.fcs = nn.ModuleList([])
        for i in range(len(kernels)):
            self.fcs.append(nn.Linear(self.d, channel))
        self.softmax = nn.Softmax(dim=0)

    def forward(self, x):
        bs, c, _, _ = x.size()
        conv_outs = []
        ### split
        for conv in self.convs:
            conv_outs.append(conv(x))
        feats = torch.stack(conv_outs, 0)  # k,bs,channel,h,w

        ### fuse
        U = sum(conv_outs)  # bs,c,h,w

        ### reduction channel
        S = U.mean(-1).mean(-1)  # bs,c
        Z = self.fc(S)  # bs,d

        ### calculate attention weight
        weights = []
        for fc in self.fcs:
            weight = fc(Z)
            weights.append(weight.view(bs, c, 1, 1))  # bs,channel
        attention_weughts = torch.stack(weights, 0)  # k,bs,channel,1,1
        attention_weughts = self.softmax(attention_weughts)  # k,bs,channel,1,1

        ### fuse
        V = (attention_weughts * feats).sum(0)
        return V


if __name__ == '__main__':
    input = torch.randn(50, 512, 7, 7)
    se = SKAttention(channel=512, reduction=8)
    output = se(input)
    print(output.shape)