代码收藏家 【python】注意力机制代码
every blog every motto: You can do more than you think.
0. 前言
梳理目前主流的注意力机制代码，目前以pytorch为例。 说明：
特征图维度的组织形式为：（batch,channel,height,width）后续增加
1. 正文
1.1 SEBlock 2017
考虑通道间的注意力之间的关系，在通道上加入注意力机制 论文：https://arxiv.org/abs/1709.01507 代码：https://github.com/hujie-frank/SENet 对于输入特征图C2，其后加上SE注意力模块
1.1.1 步骤
主要分三步：
squeeze，对空间维度进行压缩，代码上即利用全局平均池化将每个通道平均成一个值，该值具有全局感受野。（通俗理解：一个通道理解为一个大饼，多个通道就是多个大饼垒在一起。全局平均池化即将一个大饼平均成一个点，整体看，类似一个垒起来的色子） 维度变化：（2,512,8,8）-> (2,512,1,1) ==> (2,512)excitation, 利用权重学习上面各通道间的相关性，代码实现有全连接和卷积核为1的卷积操作两种方式。 维度变化：（2,512）-> (2,512//reducation)->(2,512) ==>(2,512,1,1) 说明： 该过程先降维在升维，降维倍数由reducation参数决定，降低网络的计算量，其中的激活函数增加了网络的非线性。scale: 通过上面excitation的操作输出了每个通道的重要性，在通过乘法加权操作乘以输入数据C2，从而提升重要特征，抑制不重要特征。 维度变化：（2,512,8,8）*(2,512,1,1) -> (2,512,16,16)
小结： 即输入维度为（2,512,8,8），输出维度为：（2,512,8,8） 说明： 上述步骤中的“->”表示维度变化方向，“==>”表示通过view方法改变了维度。
说明：
全连接和1 × 1的卷积效果类似，上图显示为全连接，亦可为1*1的卷积，下同，不赘述。激活函数位置见代码，下同。
1.1.3 代码
1. pytorch
class SELayer(nn.Module):
def __init__(self, channel, reduction=4):
""" SE注意力机制,输入x。输入输出特征图不变
1.squeeze: 全局池化 (batch,channel,height,width) -> (batch,channel,1,1) ==> (batch,channel)
2.excitaton: 全连接or卷积核为1的卷积(batch,channel)->(batch,channel//reduction)-> (batch,channel) ==> (batch,channel,1,1) 输出y
3.scale: 完成对通道维度上原始特征的标定 y = x*y 输出维度和输入维度相同
:param channel: 输入特征图的通道数
:param reduction: 特征图通道的降低倍数
"""
super(SELayer, self).__init__()
# 自适应全局平均池化,即，每个通道进行平均池化，使输出特征图长宽为1
self.avg_pool = nn.AdaptiveAvgPool2d(1)
# 全连接的excitation
self.fc = nn.Sequential(
nn.Linear(channel, channel // reduction),
nn.ReLU(inplace=True),
nn.Linear(channel // reduction, channel),
nn.Sigmoid()
)
# 卷积网络的excitation
# 特征图变化：
# (2,512,1,1) -> (2,512,1,1) -> (2,512,1,1)
self.fc2 = nn.Sequential(
nn.Conv2d(channel, channel // reduction, 1, bias=False),
nn.ReLU(inplace=True),
nn.Conv2d(channel // reduction, channel, 1, bias=False),
nn.Sigmoid()
)
def forward(self, x):
# (batch,channel,height,width) (2,512,8,8)
b, c, _, _ = x.size()
# 全局平均池化 (2,512,8,8) -> (2,512,1,1) -> (2,512)
y = self.avg_pool(x).view(b, c)
# (2,512) -> (2,512//reducation) -> (2,512) -> (2,512,1,1)
y = self.fc(y).view(b, c, 1, 1)
# (2,512,8,8)* (2,512,1,1) -> (2,512,8,8)
pro = x * y
return x * y
2. tensorflow/keras
# SEBlock
feature_map_shape = input_x.shape # input feature map shape
x = tf.reduce_mean(x, [1, 2]) # reduce along axis 1 and 2 ,height,width,
x = Dense(feature_map_shape[-1] / 16, activation=tf.nn.relu)(x) # (batch,channel) -> (batch,channel/16)
x = Dense(feature_map_shape[-1], activation=tf.nn.relu)(x) # (batch,channel/16) -> (batch,channel)
x = tf.multiply(input_x, x) # multiply along channel
说明：
当使用全连接时，forward中平均池 (squeeze)，
# 全局平均池化 (2,512,8,8) -> (2,512,1,1) -> (2,512)
y = self.avg_pool(x).view(b, c)
当使用1*1卷积，forward中平均池化（squeeze）,
# 全局平均池化 (2,512,8,8) -> (2,512,1,1)
y = self.avg_pool(x)
1.2 CBAM 2018
论文：https://arxiv.org/abs/1807.06521 代码：https://github.com/luuuyi/CBAM.PyTorch CBAM可以无缝集成任何CNN架构中，开销不大，早期的注意力机制一种。 实验结果表明：顺序链接比并行连接好，其中通道注意力在前优于空间注意力在前。
1.2.1 通道注意力机制
1.2.1.1 概述
通道注意力机制和上面的SEBlock类似，唯一不同的是加了一个最大池化。而后，最大池化和平均池化共用一个多层感知机（mlp）, 再将结果相加和输入特征图进行点乘传入空间注意力机制。 说明： 主要步骤省略，可参考SEBlock和下面代码中的注释。
1.2.1.2 更加清晰的理解图
1.2.1.3 代码
说明： forward中，为了方便理解，展开书写了，等价于最开始注释了的几行。
class ChannelAttention(nn.Module):
def __init__(self, in_channel, ratio=16):
""" 通道注意力机制 同最大池化和平均池化两路分别提取信息，后共用一个多层感知机mlp,再将二者结合
:param in_channel: 输入通道
:param ratio: 通道降低倍率
"""
super(ChannelAttention, self).__init__()
# 平均池化
self.avg_pool = nn.AdaptiveAvgPool2d(1)
# 最大池化
self.max_pool = nn.AdaptiveMaxPool2d(1)
# 通道先降维后恢复到原来的维数
self.fc1 = nn.Conv2d(in_channel, in_channel // ratio, 1, bias=False)
self.relu1 = nn.ReLU()
self.fc2 = nn.Conv2d(in_channel // ratio, in_channel, 1, bias=False)
self.sigmoid = nn.Sigmoid()
def forward(self, x):
# 平均池化
# avg_out = self.fc2(self.relu1(self.fc1(self.avg_pool(x))))
# 最大池化
# max_out = self.fc2(self.relu1(self.fc1(self.max_pool(x))))
# out = avg_out + max_out
# return x*self.sigmoid(out)
# 平均池化一支 (2,512,8,8) -> (2,512,1,1) -> (2,512/ration,1,1) -> (2,512,1,1)
# (2,512,8,8) -> (2,512,1,1)
avg = self.avg_pool(x)
# 多层感知机mlp (2,512,8,8) -> (2,512,1,1) -> (2,512/ration,1,1) -> (2,512,1,1)
# (2,512,1,1) -> (2,512/ratio,1,1)
avg = self.fc1(avg)
avg = self.relu1(avg)
# (2,512/ratio,1,1) -> (2,512,1,1)
avg_out = self.fc2(avg)
# 最大池化一支
# (2,512,8,8) -> (2,512,1,1)
max = self.max_pool(x)
# 多层感知机
# (2,512,1,1) -> (2,512/ratio,1,1)
max = self.fc1(max)
max = self.relu1(max)
# (2,512/ratio,1,1) -> (2,512,1,1)
max_out = self.fc2(max)
# (2,512,1,1) + (2,512,1,1) -> (2,512,1,1)
out = avg_out + max_out
return x*self.sigmoid(out)
1.2.2 空间注意力机制
1.2.2.1 概述
将通道注意力机制的的结果作为新的输入特征图，主要步骤：
输入特征图，经过最大池化,平均池化（通道维度压缩，与前面的通道注意力机制不同） 维度变化：（2,512,8，8 ） -> (2,1,8,8)将最大池化和平均池化在通道方向上合并 维度变化：（2,1,8,8）+ (2,1,8,8) -> (2,2,8,8) 3.经过卷积，通道变为1,再经过激活函数 维度变化：（2,2，8,8）-> (2,1,8,8)和输入特征图点乘 维度变化：(2,512,8,8) * (2,1,8,8) -> (2,512,8,8)
1.2.2.2 更加清晰的理解图
说明： 下面两种图例其实是一个意思，即，单通道特征图，为了看起来更加清晰用了两种颜色。
1.2.2.3 代码
class SpatialAttention(nn.Module):
def __init__(self, kernel_size=7):
""" 空间注意力机制 将通道维度通过最大池化和平均池化进行压缩，然后合并，再经过卷积和激活函数，结果和输入特征图点乘
:param kernel_size: 卷积核大小
"""
super(SpatialAttention, self).__init__()
assert kernel_size in (3, 7), 'kernel size must be 3 or 7'
padding = 3 if kernel_size == 7 else 1
self.conv1 = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False)
self.sigmoid = nn.Sigmoid()
def forward(self, x):
print('x shape', x.shape)
# (2,512,8,8) -> (2,1,8,8)
avg_out = torch.mean(x, dim=1, keepdim=True)
# (2,512,8,8) -> (2,1,8,8)
max_out, _ = torch.max(x, dim=1, keepdim=True)
# (2,1,8,8) + (2,1,8,8) -> (2,2,8,8)
cat = torch.cat([avg_out, max_out], dim=1)
# (2,2,8,8) -> (2,1,8,8)
out = self.conv1(cat)
return x * self.sigmoid(out)来源：胡侃有料
