这是目标分割系列的第7篇。前面几期把FCN、U-Net、SegNet这些结构聊得差不多了这期要聊的PSPNet是语义分割绕不开的一个“上下文聚合”里程碑。我在自己的项目里拿它当baseline用过很久也帮人调过不少基于它的分割模型所以这篇不打算只贴论文结构而是把“为什么要这么设计”“每个模块到底在算什么东西”“复现和调参时容易在哪儿翻车”一起讲透。如果你正准备搭一个语义分割模型或者已经试过FCN/DeepLab但总觉得大块区域分割得不够干净那这篇应该能帮你省不少时间。就算你只是想在面试或者组会里把PSPNet讲明白看完这篇也能应付得过来。1. 为什么分割输出总像“打满补丁”感受野与上下文这两个老问题1.1 一张图里“小目标”和“大区域”的难度完全不在一个量级先说一个大家训练分割模型时大概率遇到过的现象小目标分割得挺干净一个人、一辆车、一个交通牌边界描得很利索但一遇到整片天空、整面白墙、一大片草地输出就碎成渣像打满了补丁这儿一块那儿一块。原因说起来很直白卷积网络判断一个像素是什么类别主要看它周围的感受野里有什么。小目标本身纹理独特、边界清晰局部信息就足够区分大区域恰恰相反一片天空在不同位置可能颜色、亮度、云层分布差异很大单看局部特征模型会把它误判成“白色的墙”“浅色的路”只有把更大范围的上下文拉进来才能确认“这整片东西其实是同一种东西”。这就是语义分割里被说烂了但永远重要的问题上下文信息。PSPNet提出的时候就是奔着这个问题去的。1.2 从FCN到带空洞卷积的ResNet感受野是怎么一步步扩大的要理解PSPNet得先理解它之前的网络在感受野上吃了什么亏。FCN是最早把全连接层换成卷积层做分割的但它受限于骨干网络的步长最终特征图往往只有原图的1/32大小再上采样回去做预测。这时候一个像素对应的感受野虽然不小但特征图分辨率太低边界细节已经丢了。后来的做法是保留骨干网络比如VGG、ResNet在其顶层接解码器然后想法子让顶层特征图分辨率别降得那么狠。ResNet出来之后大家对骨干网络越来越依赖但ResNet默认的stride32对于分割来说实在太狠了。于是有了一个很常规的改造把最后几层的stride改成1用空洞卷积补偿视野。举个具体例子一个3×3卷积dilation2等效视野是7×7dilation4等效视野是15×15。这样一步步把全局stride从32拉到16甚至8特征图分辨率上去了感受野也不缩水。但问题来了即便感受野在数值上够了模型也不一定真的“用得上”这些远处的信息。普通卷积在这种长距离依赖上效率很低信息要被一层层传递才能汇合远距离像素之间的关联早就衰减得差不多了。这就是PSPNet要解决的真正痛点如何把全局信息“直接打包”送到每个像素的预测里去。2. PSPNet的正菜金字塔池化模块PPM内部拆解2.1 四个bin size分别“看”到了什么PSPNet的核心是一个叫做Pyramid Pooling Module金字塔池化模块的东西大家一般叫它PPM。它的结构看起来特别简单——用四种不同尺度的自适应平均池化把特征图分别压成1×1、2×2、3×3、6×6的小特征图然后各自过一层1×1卷积再全部上采样回原尺寸和原始特征图拼在一起。我第一眼看这个结构的时候觉得“就这么点东西”后来自己手写过一遍才明白这四个分支每一个都有它存在的意义。1×1分支相当于把整张特征图压成一个“平均值”。这一步在论文里称为全局场景级别的上下文说人话就是让网络知道“我是一整张图我大概率是户外的城市街景”这种最粗粒度的信息。它给每个像素提供了一个“整图投票”的结果。2×2分支把特征图切成4块每块算一个平均。这等于告诉网络“这块区域更偏向天空还是建筑”。3×3分支切成9块粒度更细一点能捕捉“局部区域里的大块目标分布”。6×6分支切成36块每个网格已经比较小了。它捕捉的不是全局而是把一个中等尺度的区块信息“浓缩”成更紧凑的表示。这四个分支拼在一起等于同时看了“整张图”“四分之一张图”“九分之一张图”和“三十六分之一张图”。不同大目标对上下文尺度的需求不一样有的可能需要全图信息才能判断有的只需要周围一个街区就够多尺度同时给出让网络自己去学该信谁。2.2 特征融合后通道数与分辨率的账怎么算PPM里有一个关键技术细节每个分支的池化输出后都接一个1×1卷积把通道数降到总通道数的1/4。以ResNet50输出的2048通道为例每个分支的把通道降到512。为什么非要降维这笔账其实算一下就知道。如果不降维1×1、2×2、3×3、6×6四个分支分别对应1493650张特征图。拿6×6那个分支说它输出是6×6×2048就算上采样放大回去后面跟着的卷积计算量也翻了好几倍显存直接吃紧。降维到512之后四个分支加起来是50×51225600个通道值再加上原始2048维总量控制得比较舒服。另一个细节是上采样。这里用的是双线性插值不是可学习的转置卷积。为什么因为PPM的目的是把“低分辨率的上下文统计信息”贴回原来的位置这一过程不要求引入额外的学习参数。用双线性插值参数少、训练稳定而且不产生转置卷积常见的那种棋盘格伪影。如果你手痒把这里换成反卷积大概率不会变好只会让训练更飘。最后把四个分支上采样后的特征和原始特征在通道维度上拼接得到的特征再走一层卷积做融合就得到了既保留局部细节、又带全局上下文的最终表示。这一步是整个PSPNet“力大砖飞”的枢纽。3. 整体网络怎么搭骨干改造、辅助损失与训练设置3.1 主干网络的选择与空洞卷积改造PSPNet最初用的是ResNet50和ResNet101现在复现一般也默认这两个。选ResNet当骨干有几个好处它已经足够深ImageNet预训练权重好找而且它的stage天然是金字塔式的下采样结构正好和后面接PPM的风格匹配。关键在于改造。ResNet默认输出的特征图是原图的1/32直接用会让最后分割输出太糊。常见做法是把res4和res5两个stage的stride改为1同时把res5里的3×3卷积换成空洞卷积rate一般取4或者配合rate2。这样总stride从32降到8特征图分辨率变为原来的4倍同时感受野不会因为去掉下采样而缩水。这里有个非常容易踩的坑改完stride之后如果加载的是原版ImageNet预训练权重网络结构的下采样路径变了前期统计量会有一点点不匹配。虽然大部分复现直接硬加载也没啥大问题但如果你的数据集和ImageNet差很远我建议用带空洞卷积的预训练权重或者干脆先在目标任务小数据上微调一段时间再切到分割头上训练。3.2 辅助损失一个训练期才存在的加速器PSPNet里还有个常被新手忽略的设计辅助损失auxiliary loss。在主干网络的某个中间层特征图之后额外接一个小的分类头单独算一份分割损失和主损失加权加总。原文里这个辅助损失的权重大概是0.4主损失是1.0而且辅助分支只在训练时存在推理时直接丢掉。这个设计的目的很朴素网络太深梯度从最后的分割头一路传回早期卷积层时已经很弱了。中间加一个监督信号相当于在最难传梯度的位置“接力”一下让前面的层也早点学到语义信息。我在实际使用中的感受是加上辅助损失之后训练初期loss曲线明显更平稳收敛也快一些。不过要注意辅助损失加的层不能太靠前太靠前的话监督信号过于琐碎反而可能干扰主干特征提取。一般放在res4之后比较稳妥。还有个容易忽略的事如果你加了辅助分支测试时一定要记得走不带辅助分支的推理路径。我见过不止一个人把训练代码和验证代码写成同一个forward结果测试时也把辅助分支的输出加进去最后指标莫名其妙低了一截。另外BN层在训练和推理之间的差异也很关键。如果batch size开得比较小比如只有2、4BN的均值方差估计会很不稳分割效果会明显变差。这种情况建议直接换成GroupNorm或者用同步BN把多卡上的样本合起来算统计量。属于老生常谈但在PSPNet这种大输入尺寸的模型上尤为明显——显存压力大batch size很难开大。4. PSPNet代码拆解PyTorch实现金字塔池化的关键片段4.1 核心模块PPM的搭建纸上谈兵说了那么多直接上代码比什么都直观。一个标准的PPM模块用PyTorch写起来大概是这样import torch import torch.nn as nn import torch.nn.functional as F class PyramidPoolingModule(nn.Module): def __init__(self, in_channels, bin_sizes(1, 2, 3, 6), out_channels512): super().__init__() self.branches nn.ModuleList() for bins in bin_sizes: self.branches.append(nn.Sequential( nn.AdaptiveAvgPool2d(bins), nn.Conv2d(in_channels, out_channels, kernel_size1, biasFalse), nn.BatchNorm2d(out_channels), nn.ReLU(inplaceTrue), )) def forward(self, x): h, w x.size(2), x.size(3) features [x] for branch in self.branches: y branch(x) y F.interpolate(y, size(h, w), modebilinear, align_cornersTrue) features.append(y) return torch.cat(features, dim1)几个值得注意的点第一AdaptiveAvgPool2d(bins)是自动把特征图池化成任意尺寸不需要手动算核大小。比如输入是30×60的特征图你想得到6×6的输出它会自动决定池化窗口和步长不用自己抠细节。第二1×1卷积之前没有激活函数这是常规做法。先降维再BN再ReLU顺序别搞反。第三F.interpolate上去时用的尺寸是原始特征图的h, w不是某个固定值这样输入分辨率无论多大都能跑。这个设计在PSPNet里很重要毕竟分割任务的输入尺寸经常不一样。4.2 带辅助分支的完整前向结构光一个PPM还不够要把它嵌进完整网络里。下面是我平时用来做实验的简化版PSPNet结构主干用ResNet50最后接PPM然后分割头出预测同时保留一个辅助分支import torchvision.models as models class PSPNet(nn.Module): def __init__(self, n_classes, backboneresnet50, use_auxTrue): super().__init__() self.use_aux use_aux resnet models.resnet50(pretrainedTrue) # 假设resnet已经按空洞卷积方式改造过 # 这里直接取res4和res5的输出作为特征。 self.layer0 nn.Sequential(resnet.conv1, resnet.bn1, resnet.relu, resnet.maxpool) self.layer1 resnet.layer1 self.layer2 resnet.layer2 self.layer3 resnet.layer3 # 输出为1024通道 self.layer4 resnet.layer4 # 输出为2048通道stride8 self.ppm PyramidPoolingModule(2048, bin_sizes(1, 2, 3, 6), out_channels512) self.final nn.Sequential( nn.Conv2d(2048 512 * 4, 2048, kernel_size3, padding1, biasFalse), nn.BatchNorm2d(2048), nn.ReLU(inplaceTrue), nn.Dropout2d(0.1), nn.Conv2d(2048, n_classes, kernel_size1), ) if use_aux: self.aux_head nn.Sequential( nn.Conv2d(1024, 256, kernel_size3, padding1, biasFalse), nn.BatchNorm2d(256), nn.ReLU(inplaceTrue), nn.Dropout2d(0.1), nn.Conv2d(256, n_classes, kernel_size1), ) def forward(self, x): x self.layer0(x) x self.layer1(x) x self.layer2(x) c3 self.layer3(x) c4 self.layer4(c3) p self.ppm(c4) out self.final(p) if self.training and self.use_aux: aux self.aux_head(c3) return out, aux return out注意这里self.training这个判断很关键。训练时返回两个输出主loss和辅助loss一起算推理时只返回主预测结果。还有一点代码里self.final第一层卷积用的是3×3而不是1×1目的是让拼接后的特征在融合时能参考一点邻近位置的信息。论文原文里也有类似操作能让融合后的特征更平滑。我自己试过把3×3换成1×1分割边缘会明显糙一点。实际工程里还需要把ResNet的layer4中最后一个block的stride改成1并在卷积里设置dilation否则特征图还是1/16或1/32分辨率。这块改法各家不完全一样有的只改最后两个block有的直接改整个layer4效果差异不算很大但一定要保证总stride是8左右太小了显存放不下太大了边界细节又不够。5. PSPNet的实测表现与调参经验5.1 在公开数据集上的基准成绩PSPNet提出时在公开数据集上的表现放到今天看其实已经不算顶尖但作为baseline它依然很有参考价值。以当初最常刷的两个数据集为例在Cityscapes数据集上PSPNet配合ResNet101测试集mIoU能到80%上下在PASCAL VOC 2012上多尺度融合后接近85%。具体数字每个版本和训练配置略有浮动但量级大概就是这个水平。我印象更深的是它在当时解决了一个很典型的问题大物体、大区域的连续性大幅提升。以前FCN在Cityscapes这种街景数据上墙面、道路、天空这些类别经常被分割得零零碎碎PSPNet出来后这类错误明显少了。这也直接印证了PPM的那套逻辑——把全局信息粘回每个像素大块目标自然就不容易碎。如果你要拿它做对比实验建议就用上面代码那个结构当baseline不要一开始就上多尺度测试和flip增强先把单尺度结果跑稳再逐步加trick这样你能看出每一步到底涨了多少。5.2 训练细节poly学习率、数据增强和输入的坑PSPNet原文训练时用的是poly学习率策略公式是当前学习率等于初始学习率乘以(1 - 当前迭代数/总迭代数)的0.9次方。这种方式比固定学习率或者step下降要稳很多算是分割任务里的标配了。初始学习率一般取0.01左右配合SGD和momentum0.9weight decay取1e-4。数据增强方面随机翻转、随机缩放、随机裁剪这几样基本不能少。裁剪尺寸不用贪大常见的是473或者512见方够用就行。输入分辨率对PSPNet的影响非常大这块要特别注意分辨率太低6×6那个分支能捕捉的局部上下文就少分辨率太高显存压力又顶不住。我自己的习惯是先按数据集特点定一个基础尺寸然后训练时做多尺度随机缩放比如从0.75到1.25倍之间随机采样推理时再用一个固定尺寸或原图尺寸。还有两个容易被忽略的点第一辅助分支的权重不要一上来就调太高。0.4是一个比较稳妥的默认值如果你发现主loss的值一直在高水位下不去先把辅助权重降到0.2试试。第二dropout层别忘了。final那段里我加了一个0.1的Dropout2d有人复现的时候会偷懒去掉。在数据量一般的情况下这一层能防过拟合涨点虽然不多但白送的分不要丢。另外关于测试多尺度推理是PSPNet这类模型涨点的常规套路把输入缩放到0.75、1.0、1.25三档分别预测再把softmax结果平均。这个操作一般能带来1~2个点的mIoU提升代价是推理时间翻倍。如果产品对速度有要求单尺度就够了不用硬上。6. PPM和ASPP的对照以及它在今天的“余热”6.1 多尺度感受野的两种打开方式很多人会把PSPNet和DeepLab系列放在一起比较因为DeepLab里的ASPP空洞空间金字塔池化和PPM干的是类似的事用多个并行分支抓取不同尺度的上下文。但两者的实现思路很不一样。我把它们的区别整理成一个表方便快速对照维度PPMPSPNetASPPDeepLab分支方式不同尺寸的平均池化不同rate的空洞卷积全局信息1×1分支就是全局平均池化一般单独加一个全局平均池化分支参数开销小主要是池化和1×1卷积大每个空洞卷积都带完整卷积核局部细节保留池化会损失一些细节靠拼接原始特征弥补空洞卷积保留更多细节显存占用相对低相对高训练稳定性比较稳对rate和BN比较敏感理解这个对比对实际选型很有用。如果你的显存紧、数据集不算特别大PSPNet这套更省心如果数据集很丰富且你追求极致精度DeepLab系列的天花板通常更高一点。但很难说谁全面碾压谁更多还是看场景。6.2 现在的工程里PSPNet还能用来做什么这几年Transformer分割模型、mask transformer这些新东西层出不穷动不动就刷榜。但真正落到工程上PSPNet并没有被淘汰反而在一些特定场景里很好使。我自己在几个方向见过它被用作主力模型或强力baseline中低分辨率场景的实时分割比如工业质检里把图像分成“正常区域”“缺陷区域”输入分辨率一般不会特别高PSPNet的轻量版本可以跑到接近实时的速度。遥感影像分割遥感图里建筑物、植被、水体都是大块连续区域和PSPNet擅长处理的场景极其匹配。作为新模型的对比baseline发论文或者做技术预研时PSPNet跑出来基本等同于“经典方法的下限参考”后面想验证新模块有没有用跟它比最公平。换骨干网络后的复用PPM本身不挑骨干你可以把Swin Transformer、ConvNeXt之类的新骨干和PPM拼在一起用效果通常比直接换解码器更稳。我还见过一个做法把PPM里的6×6分支改成8×8或者10×10来做遥感小目标分割。理论上网格越细局部信息越强但实践下来并非越大越好因为每个格子的区域越来越小池化的“平均”特性会让小目标被抹平。如果你也想改这个参数建议从6开始小步试别指望一次调到最优。最后分享一个实操经验如果你第一次跑PSPNet先把PPM四个分支的输出分别可视化出来看一眼。很多时候你会意外发现2×2和3×3分支上采样后的热力图长得特别像这说明两个分支学到的信息有冗余可以试试去掉一个或者调整bin size。这个细节原论文里不会写但却是你自己调模型时真正有价值的那几步。分割任务最怕的不是模型不够深而是模型看不见上下文。PPM这套方案看着简单但它在“把全局上下文带进每个像素”这件事上确实讲得明明白白。