注意力机制(SE,ECA,CBAM) Pytorch代码
2023年,3月2号,新增SKNet代码
1 SENet
SE注意力机制(Squeeze-and-Excitation Networks)
:是一种通道类型的注意力机制,就是在通道维度上增加注意力机制,主要内容是是squeeze和excitation.
就是使用另外一个新的神经网络(两个Linear层),针对通道维度的数据进行学习,获取到特征图每个通道的重要程度,然后再和原始通道数据相乘即可。
具体参考Blog:
CNN中的注意力机制
小结:
-
SENet的核心思想是通过全连接网络根据loss损失来自动学习特征权重,而不是直接根据特征通道的数值分配来判断,使有效的特征通道的权重大。
-
论文认为excitation操作中使用两个全连接层相比直接使用一个全连接层,它的好处在于,具有更多的非线性,可以更好地拟合通道间的复杂关联。
代码:
拆解步骤,forward代码写的比较细节
import torch
from torch import nn
from torchstat import stat # 查看网络参数# 定义SE注意力机制的类
class se_block(nn.Module):# 初始化, in_channel代表输入特征图的通道数, ratio代表第一个全连接下降通道的倍数def __init__(self, in_channel, ratio=4):# 继承父类初始化方法super(se_block, self).__init__()# 属性分配# 全局平均池化,输出的特征图的宽高=1self.avg_pool = nn.AdaptiveAvgPool2d(output_size=1)# 第一个全连接层将特征图的通道数下降4倍self.fc1 = nn.Linear(in_features=in_channel, out_features=in_channel//ratio, bias=False)# relu激活self.relu = nn.ReLU()# 第二个全连接层恢复通道数self.fc2 = nn.Linear(in_features=in_channel//ratio, out_features=in_channel, bias=False)# sigmoid激活函数,将权值归一化到0-1self.sigmoid = nn.Sigmoid()# 前向传播def forward(self, inputs): # inputs 代表输入特征图# 获取输入特征图的shapeb, c, h, w = inputs.shape# 全局平均池化 [b,c,h,w]==>[b,c,1,1]x = self.avg_pool(inputs)# 维度调整 [b,c,1,1]==>[b,c]x = x.view([b,c])# 第一个全连接下降通道 [b,c]==>[b,c//4] # 这里也是使用Linear层的原因,只是对Channel进行线性变换x = self.fc1(x)x = self.relu(x)# 第二个全连接上升通道 [b,c//4]==>[b,c] # 再通过Linear层恢复Channel数目x = self.fc2(x)# 对通道权重归一化处理 # 将数值转化为(0,1)之间,体现不同通道之间重要程度x = self.sigmoid(x)# 调整维度 [b,c]==>[b,c,1,1] x = x.view([b,c,1,1])# 将输入特征图和通道权重相乘outputs = x * inputsreturn outputs
结果展示:
提示:
in_channel/ratio需要大于0,否则线性层输入是0维度,没有意义,可以根据自己需求调整ratio的大小。
2 ECANet
作者表明 SENet
中的降维会给通道注意力机制带来副作用,并且捕获所有通道之间的依存关系是效率不高的,而且是不必要的。
参考Blog:
CNN中的注意力机制
代码:
详细版本:在forward中,介绍了每一步的作用
import torch
from torch import nn
import math
from torchstat import stat # 查看网络参数# 定义ECANet的类
class eca_block(nn.Module):# 初始化, in_channel代表特征图的输入通道数, b和gama代表公式中的两个系数def __init__(self, in_channel, b=1, gama=2):# 继承父类初始化super(eca_block, self).__init__()# 根据输入通道数自适应调整卷积核大小kernel_size = int(abs((math.log(in_channel, 2)+b)/gama))# 如果卷积核大小是奇数,就使用它if kernel_size % 2:kernel_size = kernel_size# 如果卷积核大小是偶数,就把它变成奇数else:kernel_size = kernel_size + 1# 卷积时,为例保证卷积前后的size不变,需要0填充的数量padding = kernel_size // 2# 全局平均池化,输出的特征图的宽高=1self.avg_pool = nn.AdaptiveAvgPool2d(output_size=1)# 1D卷积,输入和输出通道数都=1,卷积核大小是自适应的# 这个1维卷积需要好好了解一下机制,这是改进SENet的重要不同点self.conv = nn.Conv1d(in_channels=1, out_channels=1, kernel_size=kernel_size,bias=False, padding=padding)# sigmoid激活函数,权值归一化self.sigmoid = nn.Sigmoid()# 前向传播def forward(self, inputs):# 获得输入图像的shapeb, c, h, w = inputs.shape# 全局平均池化 [b,c,h,w]==>[b,c,1,1]x = self.avg_pool(inputs)# 维度调整,变成序列形式 [b,c,1,1]==>[b,1,c]x = x.view([b,1,c]) # 这是为了给一维卷积# 1D卷积 [b,1,c]==>[b,1,c]x = self.conv(x)# 权值归一化x = self.sigmoid(x)# 维度调整 [b,1,c]==>[b,c,1,1]x = x.view([b,c,1,1])# 将输入特征图和通道权重相乘[b,c,h,w]*[b,c,1,1]==>[b,c,h,w]outputs = x * inputsreturn outputs
精简版:
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchinfo import summary
import mathclass EfficientChannelAttention(nn.Module): # Efficient Channel Attention moduledef __init__(self, c, b=1, gamma=2):super(EfficientChannelAttention, self).__init__()t = int(abs((math.log(c, 2) + b) / gamma))k = t if t % 2 else t + 1self.avg_pool = nn.AdaptiveAvgPool2d(1)self.conv1 = nn.Conv1d(1, 1, kernel_size=k, padding=int(k/2), bias=False)self.sigmoid = nn.Sigmoid()def forward(self, x):x = self.avg_pool(x)# 这里可以对照上一版代码,理解每一个函数的作用x = self.conv1(x.squeeze(-1).transpose(-1, -2)).transpose(-1, -2).unsqueeze(-1)out = self.sigmoid(x)return out
效果展示:
总结:
ECANet参数更少!
3 CBAM
CBAM注意力机制是由通道注意力机制(channel)和空间注意力机制(spatial)组成。
先通道注意力,后空间注意力的顺序注意力模块!
3.1 通道注意力
输入数据,对数据分别做最大池化操作和平均池化操作(输出都是batchchannel11),然后使用SENet的方法,针对channel进行先降维后升维操作,之后将输出的两个结果相加,再使用Sigmoid得到通道权重,再之后使用View函数恢复(batchchannel11)维度,和原始数据相乘得到通道注意力结果!
通道注意力代码:
#(1)通道注意力机制
class channel_attention(nn.Module):# 初始化, in_channel代表输入特征图的通道数, ratio代表第一个全连接的通道下降倍数def __init__(self, in_channel, ratio=4):# 继承父类初始化方法super(channel_attention, self).__init__()# 全局最大池化 [b,c,h,w]==>[b,c,1,1]self.max_pool = nn.AdaptiveMaxPool2d(output_size=1)# 全局平均池化 [b,c,h,w]==>[b,c,1,1]self.avg_pool = nn.AdaptiveAvgPool2d(output_size=1)# 第一个全连接层, 通道数下降4倍self.fc1 = nn.Linear(in_features=in_channel, out_features=in_channel//ratio, bias=False)# 第二个全连接层, 恢复通道数self.fc2 = nn.Linear(in_features=in_channel//ratio, out_features=in_channel, bias=False)# relu激活函数self.relu = nn.ReLU()# sigmoid激活函数self.sigmoid = nn.Sigmoid()# 前向传播def forward(self, inputs):# 获取输入特征图的shapeb, c, h, w = inputs.shape# 输入图像做全局最大池化 [b,c,h,w]==>[b,c,1,1]max_pool = self.max_pool(inputs)# 输入图像的全局平均池化 [b,c,h,w]==>[b,c,1,1]avg_pool = self.avg_pool(inputs)# 调整池化结果的维度 [b,c,1,1]==>[b,c]max_pool = max_pool.view([b,c])avg_pool = avg_pool.view([b,c])# 第一个全连接层下降通道数 [b,c]==>[b,c//4]x_maxpool = self.fc1(max_pool)x_avgpool = self.fc1(avg_pool)# 激活函数x_maxpool = self.relu(x_maxpool)x_avgpool = self.relu(x_avgpool)# 第二个全连接层恢复通道数 [b,c//4]==>[b,c]x_maxpool = self.fc2(x_maxpool)x_avgpool = self.fc2(x_avgpool)# 将这两种池化结果相加 [b,c]==>[b,c]x = x_maxpool + x_avgpool# sigmoid函数权值归一化x = self.sigmoid(x)# 调整维度 [b,c]==>[b,c,1,1]x = x.view([b,c,1,1])# 输入特征图和通道权重相乘 [b,c,h,w]outputs = inputs * xreturn outputs
3.2 空间注意力
针对输入数据,分别选取数据中最大值所在的维度(batch1h*w),和按照维度进行数据平均操作(batch1hw),然后将两个数据做通道连接(batch2hw),使用卷积操作,将channel维度降为1,之后对结果取sigmoid,得到空间注意力权重,和原始数据相乘得到空间注意力结果。
代码:
#(2)空间注意力机制
class spatial_attention(nn.Module):# 初始化,卷积核大小为7*7def __init__(self, kernel_size=7):# 继承父类初始化方法super(spatial_attention, self).__init__()# 为了保持卷积前后的特征图shape相同,卷积时需要paddingpadding = kernel_size // 2# 7*7卷积融合通道信息 [b,2,h,w]==>[b,1,h,w]self.conv = nn.Conv2d(in_channels=2, out_channels=1, kernel_size=kernel_size,padding=padding, bias=False)# sigmoid函数self.sigmoid = nn.Sigmoid()# 前向传播def forward(self, inputs):# 在通道维度上最大池化 [b,1,h,w] keepdim保留原有深度# 返回值是在某维度的最大值和对应的索引x_maxpool, _ = torch.max(inputs, dim=1, keepdim=True)# 在通道维度上平均池化 [b,1,h,w]x_avgpool = torch.mean(inputs, dim=1, keepdim=True)# 池化后的结果在通道维度上堆叠 [b,2,h,w]x = torch.cat([x_maxpool, x_avgpool], dim=1)# 卷积融合通道信息 [b,2,h,w]==>[b,1,h,w]x = self.conv(x)# 空间权重归一化x = self.sigmoid(x)# 输入特征图和空间权重相乘outputs = inputs * xreturn outputs
3.3 CBAM
将通道注意力模块和空间注意力模块顺序串联得到CBAM模块!
代码:
class cbam(nn.Module):# 初始化,in_channel和ratio=4代表通道注意力机制的输入通道数和第一个全连接下降的通道数# kernel_size代表空间注意力机制的卷积核大小def __init__(self, in_channel, ratio=4, kernel_size=7):# 继承父类初始化方法super(cbam, self).__init__()# 实例化通道注意力机制self.channel_attention = channel_attention(in_channel=in_channel, ratio=ratio)# 实例化空间注意力机制self.spatial_attention = spatial_attention(kernel_size=kernel_size)# 前向传播def forward(self, inputs):# 先将输入图像经过通道注意力机制x = self.channel_attention(inputs)# 然后经过空间注意力机制x = self.spatial_attention(x)return x
结果:
4 展示网络层具体信息
安装包
pip install torchstat
使用
from torchstat import stat net = cbam(16)
stat(net, (16, 256, 256)) # 不需要Batch维度
5 SKNet
这是SENet的改进版,增加了多个分支,每个分支的感受野不同。
论文:https://arxiv.org/pdf/1903.06586
代码:
'''
Descripttion:
Result:
Author: Philo
Date: 2023-03-02 14:55:44
LastEditors: Philo
LastEditTime: 2023-03-02 16:01:03
'''
import torch.nn as nn
import torchclass SKConv(nn.Module):def __init__(self, in_ch, M, G, r, stride=1, L=32) -> None:super().__init__()""" ConstructorArgs:in_ch: input channel dimensionality.M: the number of branchs.G: num of convolution groups.r: the radio for compute d, the length of z.stride: stride, default 1.L: the minimum dim of the vector z in paper, default 32."""d = max(int(in_ch/r), L) # 用来进行线性层的输出通道,当输入数据In_ch很大时,用L就有点丢失数据了。self.M = Mself.in_ch = in_chself.convs = nn.ModuleList([])for i in range(M):self.convs.append(nn.Sequential(nn.Conv2d(in_ch, in_ch, kernel_size=3+i*2, stride=stride, padding = 1+i, groups=G),nn.BatchNorm2d(in_ch),nn.ReLU(False)))# print("D:", d)self.fc = nn.Linear(in_ch, d)self.fcs = nn.ModuleList([])for i in range(M):self.fcs.append(nn.Linear(d, in_ch))self.softmax = nn.Softmax(dim=1)def forward(self, x):for i, conv in enumerate(self.convs): # 第一部分,每个分支的数据进行相加,虽然这里使用的是torch.cat,但是后面又用了unsqueeze和sum进行升维和降维fea = conv(x).unsqueeze_(dim=1) # 这里在1这个地方新增了一个维度 16*1*64*256*256if i == 0:feas = feaelse:feas = torch.cat([feas, fea], dim=1) # feas.shape batch*M*in_ch*W*Hfea_U = torch.sum(feas, dim=1) # batch*in_ch*H*Wfea_s = fea_U.mean(-1).mean(-1) # Batch*in_chfea_z = self.fc(fea_s) # batch*in_ch-> batch*dfor i, fc in enumerate(self.fcs):# print(i, fea_z.shape)vector = fc(fea_z).unsqueeze_(dim=1) # batch*d->batch*in_ch->batch*1*in_ch# print(i, vector.shape)if i == 0:attention_vectors = vectorelse:attention_vectors = torch.cat([attention_vectors, vector], dim=1) # 同样的相加操作 # batch*M*in_chattention_vectors = self.softmax(attention_vectors) # 对每个分支的数据进行softmax操作attention_vectors = attention_vectors.unsqueeze(-1).unsqueeze(-1) # ->batch*M*in_ch*1*1fea_v = (feas * attention_vectors).sum(dim=1) # ->batch*in_ch*W*Hreturn fea_vif __name__ == "__main__":x = torch.randn(16, 64, 256, 256)sk = SKConv(in_ch=64, M=3, G=1, r=2)out = sk(x)print(out.shape)# in_ch 数据输入维度,M为分指数,G为Conv2d层的组数,基本设置为1,r用来进行求线性层输出通道的。
结果:
该注意力模块不改变输入数据的大小和维度!详细内容都已经在备注里啦,大家可以自己写一遍,走一遍代码!
注意力机制后期学习到再持续更新!!
参考博客:
CNN注意力机制
ECANet