做语义分割的朋友应该都有过这种体验网络跑起来之后小目标边缘模糊、大目标内部空洞提了各种模块效果就是上不去。我去年在一个自制的街景分割数据集上遇到过一模一样的问题后来引入了一个叫 Pyramid Squeeze Attention 的注意力模块才把mIoU拉上去一个台阶。这篇文章就把我在这个模块上从原理理解、代码实现到实验调优的完整过程写出来包括踩过的坑和一些配置建议给正在纠结注意力到底怎么加才有效的同行做个参考。1. 单尺度注意力为什么不够用一个分割项目里的真实教训1.1 一开始我只给了模型一个全局平均视野最早我用的是一种很常见的通道注意力思路把特征图做全局平均池化得到一个描述全局信息的向量再用两个1x1卷积把这个向量压缩、扩展生成一组通道权重最后用这组权重去重新缩放原来的特征图。这个结构本质上就是在告诉网络哪个通道更有用它把整个图像的统计信息压缩成了一条通道描述。在分类任务上这个做法确实很香因为分类的最终输出就是一个全局判断你给它全局平均池化出来的信息正好对路。但放到分割任务上问题就出来了分割需要的是逐像素的预测结果某一个像素到底属于路面还是车道线不仅取决于它自己的局部纹理还取决于它周围一大片区域的结构。全局平均池化确实给了一个全局视角但这个视角是一整张图只有一个数它把不同位置、不同尺度的信息全揉在一起了。我第一版分割网络就加了这种全局注意力测试下来mIoU提升非常有限只有零点几个点几乎可以忽略。我当时一度怀疑是不是实现写错了后来仔细看了特征图才发现问题靠近图像边缘的小目标比如远处的行人、路牌它们的有效上下文根本不是整个图像而是它们周边那几十个像素的范围。全局平均池化对这部分目标来说全局这个词太大了大到了没有意义。1.2 大目标和小目标互相打架第二个让我头疼的问题是同一个特征图上可能同时存在需要截然不同感受野的目标。拿街景数据来说一辆近处的大卡车可能要占半个画面你需要很大范围的上下文才能判断它是卡车而不是建筑物而远处一个行人可能只占几十个像素你反而不能用太大的感受野去描述它否则周围的信息一进来反而把它的特征给污染了。用一个固定的单尺度注意力模块等于让所有目标都使用同一个视野大小这在小尺度目标和大尺度目标混杂的场景里天然吃亏。大目标需要更多上下文但你给的注意力是全局平均池化出来的它更偏向图像的主体内容小目标需要保持局部特征但全局池化又把它的特征稀释掉了。这就导致网络像是一个偏科生要么大目标分得好、小目标一塌糊涂要么反过来很难两头兼顾。我后来尝试过堆更多层网络来扩大感受野效果也不好。层数加深之后梯度传播和训练速度都变差了而且单纯加深网络并不等于给网络提供了多尺度的建模能力它更多是提高了单尺度上的特征抽象程度。1.3 PSA想解决的核心问题不同尺度的上下文要不要分开看Pyramid Squeeze Attention这个名字拆开看其实就是三个关键词的组合Pyramid指的是金字塔结构Squeeze指的是压缩/挤压操作Attention指的是注意力机制。它想解决的核心问题就是上面说的两个痛点怎么让注意力既能看到全局、又能照顾局部还能在不同尺度之间做灵活组合。它的思路很直接与其用一个全局平均池化把所有信息压成一个向量不如把特征图并行地喂给多个不同大小的池化分支每个分支负责一个尺度层次。比如有的分支压成1x1代表最强的全局有的压成2x2保留一点空间细节有的压成4x4、8x8代表更大的空间分辨率。每个分支各自做一次注意力计算然后把多尺度的注意力结果融合起来再去调制原特征图。这样网络就不是拿到一个一刀切的注意力了而是有了好几个不同视野的注意力候选它们可以互相补充。一开始我以为这个结构会很重毕竟做了好几个并行分支但实际算下来参数量增加并不多因为每个分支的核心只是两个1x1卷积瓶颈在于通道压缩而不是空间分辨率。2. 模块设计拆解每一个设计选择背后的逻辑2.1 数据流拆开看金字塔、挤压、重标定三阶段整个模块的处理流程可以分成三个阶段。第一个阶段是金字塔构建把输入的 ( C \times H \times W ) 特征图分别做不同尺寸的自适应平均池化得到一组分辨率递减的特征图。第二个阶段是挤压与激励在每一个尺度的特征图上做通道注意力计算也就是先压缩成通道描述向量再生成通道权重。第三个阶段是空间对齐与融合把每个分支生成的注意力图插值回原始 ( H \times W ) 分辨率拼起来之后再用一个1x1卷积做融合最终得到一张逐像素的注意力调制图。我自己的实现里金字塔层级是 ( 1 \times 1、2 \times 2、4 \times 4、8 \times 8 ) 四档。需要说明的是这个选择是在多个实验中试出来的。最开始我只用了 ( 1 \times 1 ) 和 ( 4 \times 4 ) 两档效果有提升但不明显加到四档之后小目标和大目标的分类精度都出现了可感知的提升。我理解原因在于( 1 \times 1 ) 提供全局语义( 2 \times 2 ) 和 ( 4 \times 4 ) 覆盖中等尺度的目标区域( 8 \times 8 ) 保留更丰富的空间位置关系四个层级合起来才形成了一个相对连续的上下文光谱。2.2 池化层级的选择为什么是1×1、2×2、4×4、8×8池化层级的选择不是拍脑袋定的它直接决定了模块能感受到的上下文粒度。( 1 \times 1 ) 自适应池化等价于全局平均池化每个通道变成一个数这把整个图像的信息压缩成一个最强的全局线索。实测下来它对于这个区域整体上更像天空还是更像路面这类全局歧义的消除非常有效但代价是丧失所有空间位置信息。( 2 \times 2 ) 和 ( 4 \times 4 ) 属于过渡尺度它们在保留一定空间结构的同时压缩了大部分空间冗余。对于中等大小的目标这两个层级能提供目标周围一圈的上下文但又不会像全局池化那样把远处不相关信息带进来。( 8 \times 8 ) 保留的空间分辨率最大生成的注意力图最接近原图它对小目标、边缘区域、细长结构更友好。池化层级数量也不是越多越好。我试过加 ( 16 \times 16 ) 的层级效果几乎没有变化但计算量和显存占用明显增加。原因是 ( 16 \times 16 ) 相对 ( 8 \times 8 ) 提供的额外空间细节在主干网络里已经由前面的卷积层表达过了注意力模块再重复建模意义不大。而如果把最小层级去掉只保留 ( 2 \times 2、4 \times 4、8 \times 8 )全局上下文信息又会丢失大目标上的表现会回落。所以 ( 1 \times 1 ) 到 ( 8 \times 8 ) 这四档对于输入分辨率在 ( 256 ) 到 ( 512 ) 之间的特征图是一个比较稳的配置。2.3 压缩率、激活函数和融合方式怎么定每个金字塔分支里的挤压-激励部分有两个关键参数压缩率 ( r ) 和激活函数。压缩率 ( r ) 决定通道被压缩到多小我试过 ( r8、16、32 ) 三档。( r16 ) 的性价比最高参数量适中精度也稳定( r8 ) 精度提升不到半个点但参数量涨了不少( r32 ) 省参数可精度掉了约一个点。原因也好理解压缩率太大时中间层承载不了足够的通道区分度生成的注意力权重变得平滑失去了选择性。激活函数方面分支内部的激活用ReLU最后一层输出注意力权重用Sigmoid。Sigmoid把每个通道的调制系数限制在0到1之间这样做的好处是数值稳定特征图只会被抑制或保留不会被放大。如果你想做得更激进一点也可以用Tanh或者带缩放系数的变体让某些通道被增强但我在实验里发现Sigmoid在分割任务上表现最稳训练震荡最少。融合方式我对比过两种一种是直接把各分支的注意力图做逐元素相加另一种是拼接后过一个卷积。逐元素相加等效于让每个像素在多个尺度注意力之间取折中实现最简单但表达能力有限。拼接加卷积让网络有机会学习不同尺度之间的非线性组合关系实测下来mIoU比直接相加高0.3到0.5个点。代价是融合卷积增加了一点参数量但考虑到它带来的收益这个代价很值。3. 从零实现核心代码与集成细节3.1 一个可以跑起来的PyTorch实现完整的PyTorch实现并不复杂核心模块只有几十行代码。我先把实现贴出来然后逐段解释关键设计。import torch import torch.nn as nn import torch.nn.functional as F class PyramidSqueezeAttention(nn.Module): def __init__(self, in_channels, pyramid_levels(1, 2, 4, 8), reduction16): super().__init__() self.in_channels in_channels self.pyramid_levels pyramid_levels # 每个金字塔分支1x1卷积降维 - BN - ReLU - 1x1卷积恢复 - Sigmoid self.branches nn.ModuleList() for _ in pyramid_levels: self.branches.append(nn.Sequential( nn.Conv2d(in_channels, in_channels // reduction, kernel_size1), nn.BatchNorm2d(in_channels // reduction), nn.ReLU(inplaceTrue), nn.Conv2d(in_channels // reduction, in_channels, kernel_size1), nn.Sigmoid() )) # 多尺度注意力融合 self.fuse_conv nn.Sequential( nn.Conv2d(in_channels * len(pyramid_levels), in_channels, kernel_size1), nn.BatchNorm2d(in_channels), nn.ReLU(inplaceTrue) ) def forward(self, x): B, C, H, W x.shape attn_list [] for level, branch in zip(self.pyramid_levels, self.branches): if level 1: pooled x else: pooled F.adaptive_avg_pool2d(x, output_sizelevel) attn branch(pooled) attn F.interpolate(attn, size(H, W), modebilinear, align_cornersFalse) attn_list.append(attn) attn_cat torch.cat(attn_list, dim1) fused_attn self.fuse_conv(attn_cat) return x * fused_attn这个实现有几个值得注意的点。第一当 level 为 1 时我直接复用原特征图而不是先做一次池化因为自适应池化到 1x1 再插值回来跟直接对原图做全局平均池化再广播是等价的省掉一次多余的插值运算。第二每个分支内部先降维后升维是标准的瓶颈结构用来控制参数量。第三插值统一用双线性上采样并且不推荐用最近邻因为最近邻会把池化产生的块状结构直接带到注意力图里输出容易出现格子状伪影这个后面踩坑部分会详细说。3.2 集成到分割网络的位置与通道设置模块实现之后集成位置非常关键。我一开始图省事只在主干网络的最后一层输出上接了一个PSA模块效果提升有限。后来我改成在每两个下采样阶段之后的特征图上都接入模块也就是分别在四分之一分辨率、八分之一分辨率和十六分之一分辨率的特征层上各放一个PSA模块效果才真正拉开差距。原因在于低层特征分辨率高、包含细致的空间结构高层特征分辨率低、包含强语义信息。不同层需要的注意力粒度天然不同你只在最高层做多尺度注意力对低层小目标的帮助非常微弱。我建议至少放到两个不同分辨率的层上如果显存允许放三个效果更好。通道设置上模块的输入通道要和所在特征层的通道数完全一致。除了第一层PSA内部不做额外的通道变换输入输出通道保持一致这样可以用残差的方式叠加在主干特征后面集成成本很低。如果你想让它更轻可以把分支里的瓶颈通道从 in_channels // 16 改成 in_channels // 32但我在前面的实验对比里说过这样会牺牲约一个点的精度需要自己权衡。3.3 初始化、学习率和训练设置训练配置上我推荐把新加的模块初始化成恒等映射的近似状态也就是让融合卷积的初始输出尽可能接近1。不过实现里我用了Sigmoid和ReLU的组合想精确做到恒等比较麻烦。一个简单有效的替代方案是先把PSA模块的最后一层1x1卷积的偏置初始化为正值让初始的注意力图整体偏高这样模型一开始不会因为新模块的随机初始化而剧烈抖动。学习率方面我不建议给PSA单独设置一个过高的学习率。它内部的1x1卷积参数量不大正常训练就能收敛得很好。如果主干网络用了预训练权重可以把主干的学习率调成新模块的0.1倍这样新模块可以在前期快速学习同时保持主干特征的基本稳定。我用的是单卡训练batch size保持在8左右配合Poly学习率策略整体训练过程非常平稳。如果你的batch size只能开到2到4就需要小心BN的稳定性这个在踩坑部分会展开。4. 实测效果我用三个分辨率的特征图做了对比实验4.1 基线对比mIoU确实涨了为了验证模块效果我在一个自制的街景分割数据集上做了对比实验。数据集大约有一万两千张图像覆盖城市、乡村、高速公路三种场景共标注了十九个类别。主干网络用的是带预训练权重的轻量级分割网络在第三、第四、第五个下采样阶段之后各接一个PSA模块。下表是几组关键配置的实测数据所有实验都在相同的训练轮数和学习率策略下进行。配置mIoU类别平均Acc参数量增量单帧推理耗时基线无PSA68.2%75.4%-82ms只在最高层加PSA69.1%76.2%2.6M88ms三个分辨率层都加PSA71.5%78.3%6.1M98ms三个分辨率层加PSA且reduction871.8%78.6%11.2M110ms从结果可以看出三层都加PSA比只在最高层加效果好非常多提升从0.9个点变成了3.3个点。这个差距充分说明多尺度注意力是一个需要层层参与的事情只在最后一层做等于让高层来帮低层擦屁股远不如让每一层直接看到自己想要的上下文。reduction从16降到8只多带来了0.3个点的提升但参数增量接近翻倍推理也慢了12毫秒我最终没有采用这个配置。4.2 收敛曲线和损失变化前期更快后期更稳训练过程中我重点观察了每个epoch结束时的mIoU变化。带PSA的模型在前二十个epoch内mIoU爬升明显比基线快最早五个epoch就能拉开超过一个点的差距这说明多尺度注意力确实给网络提供了更强的早期学习信号。到了训练后期PSA模型的收敛曲线也更平稳。基线模型在最后二十个epoch里mIoU会出现大约0.5个点左右的小幅波动而PSA模型的波动基本在0.2个点以内。我推测原因在于多尺度注意力提供了冗余的上下文路径某个分支的过拟合不会直接导致整体预测崩坏相当于加了某种隐式的干扰抑制。从损失函数上看交叉熵损失整体下降趋势一致但PSA模型最终损失值低了约0.15。这里想提醒一点注意力模块并不是万能的它不会让你的网络从欠拟合变成过拟合更不会拯救一个本身设计就有严重问题的网络。如果你的基线网络连基础分割都做不好先把主干和损失函数调好再考虑加注意力。4.3 显存和推理速度代价很多担心模块太重的同行会先问开销。我用一张12GB显存的显卡做测试输入分辨率 ( 512 \times 512 )主干网络加上两层PSA时峰值显存大约增加了600MB。增加最多的是三层PSA且reduction8的配置额外显存接近1.2GB。对于大多数实验室里的中高端显卡这个开销完全可控。推理耗时方面PSA的额外开销主要来自两个部分一次多尺度自适应池化和多次双线性插值。自适应池化本身非常快瓶颈在把4个尺度的注意力图插值回原分辨率的过程。我在前面表格里写的是所有操作串行执行的耗时实际工程部署时你可以把四个分支的插值并行计算把耗时压到90毫秒以内。另外要注意如果用TensorRT这类推理框架导出模型双线性插值算子的实现差异会导致运行时间和精度出现变化建议导出后针对目标设备重新做精度校准。5. 验证过程中踩过的三个坑5.1 上采样方式不对模型输出出现棋盘格纹理第一次把PSA集成到完整分割网络里训练到一半我注意到预测图上有明显的棋盘格状纹理尤其在地面和墙面这些大面积平坦区域特别扎眼。一开始我怀疑是转置卷积的问题但我的网络里没有转置卷积后来单独打印PSA模块输出的注意力图发现 ( 2 \times 2 ) 和 ( 4 \times 4 ) 分支插值回原分辨率之后出现了明显的块状边界。问题根源在于我最初为了省事插值用了最近邻模式。最近邻插值不会在两个像素之间生成平滑过渡池化后的低分辨率特征被放大后网格边界被原样保留下来最终和原特征图相乘后就在输出里形成了棋盘格。把 mode 从 nearest 改成 bilinear同时把 align_corners 设为 False 之后这个问题立刻消失预测图变得平滑自然。这是一个很小但非常典型的细节坑建议大家在检查注意力模块输出时第一件事就是看它生成的注意力图是否平滑。5.2 特征图已经很小了金字塔层级还在硬撑在另一个实验里我把PSA放到了主干网络最深层那个位置的特征图分辨率只有 ( 16 \times 16 )。问题来了金字塔层级里的 ( 8 \times 8 ) 分支做自适应池化后特征图还剩 ( 8 \times 8 )信息压缩不算严重但 ( 2 \times 2 ) 和 ( 1 \times 1 ) 分支会把整个 ( 16 \times 16 ) 的特征图压成几个像素几乎把空间结构全部抹掉。插值回 ( 16 \times 16 ) 之后只能恢复出极其模糊的注意力图对原特征的调制作用约等于没有。更极端的情况是如果你的最深特征图只有 ( 8 \times 8 )那么 ( 4 \times 4 ) 以上的池化层级全都在做用稀少的信息做粗糙的注意力不仅没用还浪费算力。我后来针对深层特征专门设置了一个规则当特征图边长小于等于16时把金字塔层级降级为 ( (1, 2, 4) )小于等于8时进一步降为 ( (1, 2) )。这个调整避免了无效计算也让深层的注意力图保持了一定的空间细节最终在验证集上mIoU反而小幅提升了0.2个点。5.3 小batch下BN统计量被注意力放大第三个坑是在batch size比较小的训练场景里踩到的。我当时为了快速验证想法把batch size调成了2结果发现PSA模型训练曲线抖动非常严重loss在几个epoch内不降反升。单独看主干网络时没有这个问题加上PSA后才出现。原因是PSA每个分支内部都有BatchNorm而且它位于两个1x1卷积之间。当batch size很小的时候BN统计的均值和方差本身就非常不稳定再加上注意力图作为乘法因子直接作用到特征图上等于把BN的抖动放大了好几倍。解决办法有两个一个是在分支内部把BatchNorm换成GroupNormGroupNorm对batch size不敏感另一个是把分支里的BN去掉直接用1x1卷积加ReLU。我两个都试了去掉BN对最终精度几乎没有影响但训练稳定性明显提升。如果你的显存只允许小batch训练建议直接不加分支内部的BN。6. 什么时候用PSA什么时候该选别的方案6.1 最适用的场景多尺度目标混合的密集预测任务从我的实验经验来看PSA最适用的场景就是目标尺度跨度比较大的密集预测任务典型代表是街景分割、遥感图像分割、医疗影像中的病灶区域分割。这些任务有一个共同特点一张图里既有大面积的背景区域天空、道路、农田又有大量细小的目标行人、车辆、小病灶它们对上下文的尺度需求截然不同。在这些场景里PSA的提升主要集中在中型目标和小目标类别上。遥感图像里的建筑物边缘、街景里的远处行人精度提升格外明显。大目标和背景类别反而提升不明显因为它们本来就被主干网络的深层特征表达得不错了。6.2 和ASPP、Non-local注意力相比的取舍很多人会问PSA和ASPP有什么区别。ASPP通过在多个膨胀率的空洞卷积上提取多尺度特征它本质上是增强了特征本身的表达能力PSA做的事情不一样它是生成一个多尺度的注意力权重去调制已有特征。两者并不冲突实际上我也试过在PSA之后继续接ASPPmIoU又涨了0.4个点。相比Non-local注意力这种计算整张特征图两两相似度的方法PSA的显存和计算开销要小一个数量级。Non-local注意力在 ( 512 \times 512 ) 输入上做自注意力需要计算 ( 262144 \times 262144 ) 的相似度矩阵大多数显卡根本扛不住。PSA没有这个顾虑因为它是在池化后的低分辨率特征上做注意力然后插值回去。如果你追求接近Non-local的效果但又没有多卡或大显存的条件PSA是一个很现实的替代品。6.3 什么时候别硬上PSA有几个场景我建议不要硬上PSA。如果你的任务本身对全局上下文不敏感比如只需要识别图像中央一个固定大小的目标单尺度注意力或干脆不加注意力可能更合适。如果你的网络结构已经非常轻量比如移动端实时分割每加一个模块都要谨慎评估延时预算这时候优先考虑把PSA只放在最深层一级。另外样本量极小的情况下PSA也容易出问题。我试过一个只有两千张图像的小数据集PSA模型在训练集上表现很好验证集上反而比基线差了一个点。多尺度注意力引入了额外的表达能力在小数据集上更容易过拟合。这种情况下可以考虑把金字塔层级减少到两档或者直接放弃这个模块把预算花在数据增强上更划算。在我实际项目里的体会是Pyramid Squeeze Attention的价值不在于它有多复杂的理论而在于它用很小的开销把多尺度上下文这个概念变成了一种可以灵活叠加在任意特征层上的工程化模块。做完这一轮实验之后我形成了一条日常使用的经验清单先确认特征图分辨率会不会把金字塔层级压废再决定加在哪几层最后检查小batch下BN是否抖动。沿着这个流程走一遍大部分注意力模块效果不佳的问题其实在配置阶段就已经能避开了。