SKAttention选择核注意力

标题：SKAttention

期刊：IEEE2019

代码： https://github.com/implus/SKNet

简介：

• 动机:增大感受野来提升性能、多尺度信息聚合方式
• 解决的问题：自适应调整感受野大小
• 创新性:提出选择性内核（SK）卷积softmax来进行自适应选择

模型结构

模型代码

import numpy as np
import torch
from torch import nn
from torch.nn import init
from collections import OrderedDict# Selective Kernel Attention
class SKAttention(nn.Module):def __init__(self, channel=512, kernels=[1, 3, 5, 7], reduction=16, group=1, L=32):super().__init__()# 中间维度d的计算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激活函数('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))# 注意力权重归一化（沿分支维度softmax）self.softmax = nn.Softmax(dim=0)def forward(self, x):# 输入x形状: [B, C, H, W]bs, c, _, _ = x.size() # 获取输入的batch_size, 通道数, 高度, 宽度conv_outs = []### Split阶段：多分支特征提取for conv in self.convs:conv_outs.append(conv(x)) # 每个分支输出: [B, C, H, W]feats = torch.stack(conv_outs, 0)  # 堆叠后形状: [K, B, C, H, W]（K是kernel数量）### Fuse阶段：特征融合U = sum(conv_outs) # 逐元素相加 → [B, C, H, W]### Channel Reduction：通道压缩S = U.mean(-1).mean(-1)  # 空间全局平均池化 → [B, C，1，1]Z = self.fc(S)   # 全连接层降维 → [B, d]（d=self.d）### 计算注意力权重weights = []for fc in self.fcs: #  每个kernel对应一个全连接层weight = fc(Z) # 全连接层输出 → [B, C]weights.append(weight.view(bs, c, 1, 1))  # 调整形状 → [B, C, 1, 1]attention_weughts = torch.stack(weights, 0)   # 堆叠 → [K, B, C, 1, 1]attention_weughts = self.softmax(attention_weughts)  # 沿K维度softmax归一化### fuseV = (attention_weughts * feats).sum(0) # 加权求和 → [B, C, H, W]return Vif __name__ == '__main__':input = torch.rand(1,64,256,256).cuda()model = SKAttention(channel=64, reduction=8).cuda()output = model (input)print('input_size:', input.size())print('output_size:', output.size())print("最大内存占用:", torch.cuda.max_memory_allocated() // 1024 // 1024, "MB")