import numpy as np
import torch
from torch import nn
from torch.nn import init
from collections import OrderedDict
# selective kernel attention
# 多个卷积核的通道注意力
# 方法出处 2019 CVPR 《Selective Kernel Networks》
class SKAttention(nn.Module):
# 初始化层
def __init__(self, channel=512, kernels=[1, 3, 5, 7], reduction=16, group=1, L=32):
# 所有继承于nn.Module的模型都要写这句话
super(SKAttention, self).__init__()
# 隐藏层的缩减因子
self.d = max(L, channel // reduction)
# 不同大小的卷积核
# nn.ModuleList类似于Python中的list
# 只不过在PyTorch中
# ModuleList中存储的是神经网络层
self.convs = nn.ModuleList([])
# 遍历不同大小的卷积核
for k in kernels:
# 添加不同大小卷积核的卷积层
self.convs.append(
# nn.Sequential可以理解为不同网络层的一个组合器
# 这个OrderedDict的作用是可以为加入到Sequential的每一个神经网络层进行命名
# 这样我们可以通过名字找到相应的神经网络层
nn.Sequential(OrderedDict([
# Conv2d卷积层的输入是[N,C-in,H,W]
# 表示[批大小,输入的通道数,输入的高,输入的宽]
# group表示分组卷积的组数
# 以group=2为例
# 将输入的通道数除以2,也就是将输入的数据一分为2
# 然后对于每一部分进行卷积
# 最后将结果进行连接
# 卷积层输出的宽,高计算公式
# 输出的宽=(输入的宽+2*padding-kernel_size)/stride(卷积步幅,默认为1)+1
# 所以最后不同卷积核输出的特征图的宽和高都是相同的
# 这样后面融合的步骤才能将这些不同卷积核输出的特征图进行相加
('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))
# stack表示张量的堆叠,会在指定的维度上创建新的维度
# cat表示张量的拼接,在指定的维度上进行张量连接
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
# 计算不同通道之间的注意力得分
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
# 论文所述的不同通道注意力融合的操作
# 将不同卷积核输出的特征图乘以相应的权重
# 最后再将这些特征图进行求和
V = (attention_weughts * feats).sum(0)
return V
if __name__ == '__main__':
# 输入,可以想象成特征图
# 批大小50
# 通道512
# 宽,高7*7
input = torch.randn(50, 512, 7, 7)
se = SKAttention(channel=512, reduction=8)
output = se(input)
print(output.shape)
2019 CVPR 《Selective Kernel Networks》 PyTorch实现
猜你喜欢
转载自blog.csdn.net/Talantfuck/article/details/124567105
今日推荐
周排行