04 - CoordAttention模块
论文《Coordinate Attention for Efficient Mobile Network Design》
1、作用
Coordinate Attention提出了一种新的注意力机制，用于在移动网络中嵌入位置信息到通道注意力中。这种方法不仅关注“哪些通道是重要的”，而且关注“在哪里”关注，通过更精细地控制空间选择性注意力图的生成，进一步提升模型性能。
2、机制
1、坐标信息嵌入：
与传统的通道注意力通过2D全局池化将特征张量转换为单一特征向量不同，Coordinate Attention将通道注意力分解为两个1D特征编码过程，分别沿两个空间方向聚合特征。这种方法能够捕捉沿一个空间方向的长程依赖性，同时保留沿另一个空间方向的精确位置信息。
2、坐标注意力生成：
将沿垂直和水平方向聚合的特征图编码成一对方向感知和位置敏感的注意力图，这两个注意力图被互补地应用到输入特征图上，增强了对兴趣对象的表示。
3、独特优势
1、方向感知和位置敏感：
Coordinate Attention通过生成方向感知和位置敏感的注意力图，使模型能够更准确地定位和识别兴趣对象。这种注意力图能够精确地高亮兴趣区域，提升了模型对空间结构的理解能力。
2、灵活性和轻量级：
Coordinate Attention的设计简洁而高效，可以轻松嵌入到经典的移动网络结构中，如MobileNetV2、MobileNeXt和EfficientNet，几乎不增加计算开销，适用于计算资源受限的环境。
3、跨任务性能提升：
Coordinate Attention不仅在ImageNet分类任务上有效，更在下游任务如对象检测和语义分割上展现出更好的性能。这证明了其对于捕捉关键信息的能力，尤其在需要密集预测的任务中表现出色。
4、代码
import torch
import torch.nn as nn
import torch.nn.functional as F# 定义h_sigmoid激活函数，这是一种硬Sigmoid函数
class h_sigmoid(nn.Module):def __init__(self, inplace=True):super(h_sigmoid, self).__init__()self.relu = nn.ReLU6(inplace=inplace) # 使用ReLU6实现def forward(self, x):return self.relu(x + 3) / 6 # 公式为ReLU6(x+3)/6，模拟Sigmoid激活函数# 定义h_swish激活函数，这是基于h_sigmoid的Swish函数变体
class h_swish(nn.Module):def __init__(self, inplace=True):super(h_swish, self).__init__()self.sigmoid = h_sigmoid(inplace=inplace) # 使用上面定义的h_sigmoiddef forward(self, x):return x * self.sigmoid(x) # 公式为x * h_sigmoid(x)# 定义Coordinate Attention模块
class CoordAtt(nn.Module):def __init__(self, inp, oup, reduction=32):super(CoordAtt, self).__init__()# 定义水平和垂直方向的自适应平均池化self.pool_h = nn.AdaptiveAvgPool2d((None, 1)) # 水平方向self.pool_w = nn.AdaptiveAvgPool2d((1, None)) # 垂直方向mip = max(8, inp // reduction) # 计算中间层的通道数# 1x1卷积用于降维self.conv1 = nn.Conv2d(inp, mip, kernel_size=1, stride=1, padding=0)self.bn1 = nn.BatchNorm2d(mip) # 批归一化self.act = h_swish() # 激活函数# 两个1x1卷积，分别对应水平和垂直方向self.conv_h = nn.Conv2d(mip, oup, kernel_size=1, stride=1, padding=0)self.conv_w = nn.Conv2d(mip, oup, kernel_size=1, stride=1, padding=0)def forward(self, x):identity = x # 保存输入作为残差连接n, c, h, w = x.size() # 获取输入的尺寸x_h = self.pool_h(x) # 水平方向池化x_w = self.pool_w(x).permute(0, 1, 3, 2) # 垂直方向池化并交换维度以适应拼接y = torch.cat([x_h, x_w], dim=2) # 拼接水平和垂直方向的特征y = self.conv1(y) # 通过1x1卷积降维y = self.bn1(y) # 批归一化y = self.act(y) # 激活函数x_h, x_w = torch.split(y, [h, w], dim=2) # 将特征拆分回水平和垂直方向x_w = x_w.permute(0, 1, 3, 2) # 恢复x_w的原始维度a_h = self.conv_h(x_h).sigmoid() # 通过1x1卷积并应用Sigmoid获取水平方向的注意力权重a_w = self.conv_w(x_w).sigmoid() # 通过1x1卷积并应用Sigmoid获取垂直方向的注意力权重out = identity * a_w * a_h # 应用注意力权重到输入特征，并与残差连接相乘return out # 返回输出# 示例使用
if __name__ == '__main__':block = CoordAtt(64, 64) # 实例化Coordinate Attention模块input = torch.rand(1, 64, 64, 64) # 创建一个随机输入output = block(input) # 通过模块处理输入print(output.shape()) # 打印输入和输出的尺寸