当前位置：首页 > news >正文

22 - PSA模块

news 2025/6/15 6:35:44

论文《EPSANet: An Efficient Pyramid Squeeze Attention Block on Convolutional Neural Network》

1、作用

EPSANet通过引入高效的金字塔挤压注意力（Pyramid Squeeze Attention, PSA）模块，显著提升了深度卷积神经网络在图像分类、对象检测和实例分割等计算机视觉任务中的性能。通过在ResNet的瓶颈块中替换3x3卷积为PSA模块，EPSANet能够在不增加显著计算负担的情况下，提供更丰富的多尺度特征表示和更有效的通道空间交互。

2、机制

1、金字塔挤压注意力模块（PSA）：

通过利用多尺度金字塔卷积结构整合输入特征图的信息，并通过挤压输入张量的通道维度来有效地从每个通道的特征图中提取不同尺度的空间信息。此外，通过提取多尺度特征图的通道注意力权重并使用Softmax操作重新校准对应通道的注意力权重，建立了长程通道依赖性。

2、高效金字塔挤压注意力（EPSA）块：

将PSA模块替换到ResNet的瓶颈块中，获取了名为EPSA的新型表示块。EPSA块易于作为即插即用组件添加到现有的骨干网络中，并且能够在模型性能上获得显著提升。

3、EPSANet架构：

通过堆叠ResNet风格的EPSA块，开发了简单且高效的EPSANet骨干架构。EPSANet通过提出的PSA模块，为各种计算机视觉任务提供了更强的多尺度表示能力，并能够适应性地重新校准跨维度的通道注意力权重。

3、独特优势

1、性能提升：

相比于SENet-50，EPSANet在ImageNet数据集上的Top-1准确率提高了1.93%，在MS COCO数据集上，对象检测的box AP提高了2.7个百分点，实例分割的mask AP提高了1.7个百分点。

2、计算效率：

EPSANet模型尺寸小于高效SR方法，如IMDN，同时在内存使用上更为高效。

3、动态空间调制：

EPSANet通过动态空间调制有效聚合了特征，提升了重建性能，同时保持了低计算和存储成本。

4、有效整合局部和非局部特征：

EPSANet通过PSA层和卷积通道混合器（CCM）的结合，有效整合了局部和非局部特征信息，实现了更精确的图像超分辨率重建。

4、代码

import torch
import torch.nn as nn
import torch.nn.functional as F
# 定义PSA模块
class PSA(nn.Module):def __init__(self, channel=512, reduction=4, S=4):super(PSA, self).__init__()self.S = S  # 尺度的数量，用于控制多尺度处理的维度# 定义不同尺度的卷积层self.convs = nn.ModuleList([nn.Conv2d(channel // S, channel // S, kernel_size=2 * (i + 1) + 1, padding=i + 1)for i in range(S)])# 定义每个尺度对应的SE模块self.se_blocks = nn.ModuleList([nn.Sequential(nn.AdaptiveAvgPool2d(1), # 自适应平均池化到1x1nn.Conv2d(channel // S, channel // (S * reduction), kernel_size=1, bias=False), # 减少通道数nn.ReLU(inplace=True),# ReLU激活函数nn.Conv2d(channel // (S * reduction), channel // S, kernel_size=1, bias=False), # 恢复通道数nn.Sigmoid()# Sigmoid激活函数，输出通道注意力权重) for i in range(S)])self.softmax = nn.Softmax(dim=1)  # 对每个位置的尺度权重进行归一化def forward(self, x):b, c, h, w = x.size()# 将输入在通道维度上分割为S份，对应不同的尺度SPC_out = x.view(b, self.S, c // self.S, h, w)# 对每个尺度的特征应用对应的卷积层conv_out = []for idx, conv in enumerate(self.convs):conv_out.append(conv(SPC_out[:, idx, :, :, :]))SPC_out = torch.stack(conv_out, dim=1)# 对每个尺度的特征应用对应的SE模块，获得通道注意力权重se_out = [se(SPC_out[:, idx, :, :, :]) for idx, se in enumerate(self.se_blocks)]SE_out = torch.stack(se_out, dim=1)SE_out = SE_out.expand(-1, -1, -1, h, w)  # 扩展以匹配SPC_out的尺寸# 应用Softmax归一化注意力权重softmax_out = self.softmax(SE_out)# 应用注意力权重并合并多尺度特征PSA_out = SPC_out * softmax_outPSA_out = torch.sum(PSA_out, dim=1)  # 沿尺度维度合并特征return PSA_outif __name__ == '__main__':# 测试PSA模块input = torch.randn(3, 512, 64, 64)# 创建一个随机输入psa = PSA(channel=512, reduction=4, S=4)# 实例化PSA模块output = psa(input)# 前向传播print(output.shape)