PSMNet立体匹配实战：SPP上下文聚合与3D代价体训练全解析

做双目立体匹配的估计前几年我最常用的还是传统局部算法。直到某次项目里遇到一个头疼问题——视差图边缘撕裂感极强大片天空和玻璃区域出现雪花噪点怎么调后处理参数都压不下去。后来把主干换成PSMNet网络Pyramid Stereo Matching Network用空间金字塔池化SPP聚合多尺度上下文再用3D卷积在代价体上做全局推理预测结果一下子稳定了很多。这篇博文就把我对这个网络的理解、训练细节和踩坑记录都摊开来讲特别适合刚上手双目深度估计、想复现经典模型但又被3D代价体显存吓退的开发者。PSMNet是深度学习立体匹配里绕不开的一座桥。它不玩花活思路非常朴素第一在提取特征阶段就把多尺度上下文信息灌进去第二把匹配代价从2D平面升级成3D代价体用编码器-解码器结构做全局聚合。这两步直接拉高了当时立体匹配在公开数据集上的精度上限也带动了后续一堆3D代价体的工作。我自己在多个数据集上复现和微调过这套结构下面按“为什么这样设计——网络中到底在做什么——训练怎么落地——容易踩哪些坑”的顺序来说。1. 立体匹配做到一半卡在“视差不光滑”上——这是PSMNet要解决的事1.1 从SGM到CNN留下的那个坑传统方法里SGM半全局匹配经常被当作基线。它的思路是沿着多个方向做代价聚合本质上是在二维图像平面上传播信息。SGM的优点是快、内存占用低缺点也很明显在弱纹理、反光、遮挡区域聚合路径覆盖有限容易产生水平条纹和深度断裂。我当时在某项目里跑SGM输出在物体边界处经常出现“拖尾”就像视差被拉了一层毛边。换成CNN做立体匹配早期方案大多只做了两件事用一个共享权重的卷积网络提取左右图特征然后在特征层面计算匹配代价再通过2D卷积或简单的后处理回归视差。这类方法局部感很强因为2D卷积的感受野再大也是有限的遇到大块无纹理区域时匹配代价本身就是“聋子的耳朵”——每个像素都找不到唯一正确对应点网络根本没有足够上下文的线索去推断。1.2 上下文信息才是立体匹配的大命题理解PSMNet之前要先明白立体匹配里一个核心矛盾局部纹理信息决定“能不能匹配”全局上下文决定“匹配得对不对”。以天空为例左右视图里天空区域的像素几乎一模一样局部窗口内怎么算相似度都接近满分但这并不能告诉你视差是多少。必须借助上下文——比如天空上方是远处的山体、下方是地面通过场景分层关系才能猜出天空应该在大约哪个视差范围。再比如细长结构或黑白条纹局部匹配很容易周期性错位只有看到更大范围的其他结构才能打破歧义。所以一个立体匹配网络要在强纹理区域保留细节又要在弱纹理区域具备全局推理能力就必须有一种机制把多尺度上下文显式地送进匹配模块。传统方法里这是用多尺度的代价聚合来实现的PSMNet则用了一个更直接的方式在特征层面做空间金字塔池化。1.3 PSMNet的破局点两条腿走路的网络结构PSMNet的完整流程可以概括成三条链特征抽取链、代价体构建链、代价体聚合链。特征抽取链由基础CNN加SPP模块组成输出的是同时保留局部细节和全局线索的特征图代价体构建链把左右视图特征沿视差方向做位移拼接形成一个“高为视差范围”的3D体代价体聚合链用一套3D沙漏网络对代价体做编码-解码最后softmax之后通过加权求和回归出稠密视差。这套设计最关键的地方就是“在不牺牲局部细节的前提下把上下文直接编码进代价体”。后面出现的许多工作——GwcNet、AANet、RAFT-Stereo——其实都是在这条思路上做不同取舍。理解了PSMNet后面看那些变体就会轻松很多。2. SPP聚合模块给代价体装上“广角镜头”2.1 为什么局部卷积当不了全局聚合的主力你先想一个问题把ResNet最后一个block的特征图拿出来它的感受野理论上已经覆盖了整张图那是不是就不需要额外的全局模块了理论上是实际上不是。CNN的感受野是渐进扩大的虽然理论上能覆盖很大区域但特征图里真正有效的全局信息被层层卷积淡化了特别是对小物体和远处细节来说深层特征图分辨率也不够。SPP模块的思路简单粗暴在特征图上按不同大小的网格做池化把整个图分别压缩成几个带全局信息的特征向量再上采样回去拼接。这样既保留了原始特征图的局部细节又额外提供了“整个区域大概长什么样”的结构先验。PSMNet里用的就是经典的4层金字塔。2.2 四层金字塔池化从64到8下采样的视野扩展过程以输入图像裁剪为256×512为例经过特征提取后特征图是64×128处理前先过几个瓶颈残差块通道数为32。SPP模块在这张特征图上并行做4个不同尺寸的自适应池化输出空间尺寸分别是64×64、32×32、16×16、8×8。严格来说论文里的池化窗口是固定大小的但实际PyTorch实现多用AdaptiveAvgPool2d把特征图强制缩放到目标尺寸效果是一样的。每一层池化之后接着做两个卷积操作通常内核为3×3目的是让池化产出的一小块“区域概要”在进入融合环节之前先做一层特征提炼。最后把所有尺度的输出都用双线性插值上采样回64×128沿通道轴拼起来。你可能会问为什么不是直接用膨胀卷积扩大感受野膨胀卷积在前视场重建里确实常用但它对显存压力大而且几个不同膨胀率的卷积叠加对远处遮挡区域的整体结构建模能力不如池化来得直接。池化的好处是视野硬性全覆盖把图像分成64×64、32×32这样一个个区域网络能学着利用“这个区块代表什么物体”来推断视差。2.3 concat前的1x1卷积降维与融合的双重任务原本特征图是32通道4个金字塔分支各自输出也是32通道拼接后变成160通道。直接把这个160通道的大特征图拿去做代价体构建理论上也行但计算量会线性上涨。1×1卷积在这里承担两个任务第一个任务是通道降维。转到32通道特征图又回到了一个让后续3D卷积不至于爆炸的通道数。第二个任务是跨尺度融合。4个尺度的信息在通道维度上并排如果不做任何融合后续每个通道都只认一个尺度效果会大打折扣1×1卷积让每个输出通道都能看到拼接后所有通道的信息等于在进入匹配阶段之前让不同尺度先“打招呼”。我在实际对比实验里试过跳过SPP直接把原始特征图拿去构建代价体D1-all误差大概会涨1.2到1.5个百分点而且弱纹理区域的视差预测明显保守总是一块一块平整得像“贴纸”一样。加上SPP之后虽然参数略多但视觉上的整体感好了非常多。实测小细节如果有显存余量SPP的池化层数可以加到5到6层比如补一个128×128的浅层尺度对保留细小物体边缘有一定帮助但收益递减明显4层是最稳妥的默认配置。3. 3D代价体聚合用沙漏网络把“匹配成本”编成“视差概率”3.1 代价体怎么构建才算稳concat方式与位移细节拿到了左右图各自经过SPP的特征图f_l和f_r代价体构建就是把两个特征图按视差方向堆叠起来。假设最大视差为D通常取D192那么在视差d处把f_r整体向左平移d个像素再和f_l沿通道维拼接得到的就是一个形状为(B, 64, D, H/4, W/4)的5D张量严格说是4Dbatch和通道先不看。这里有个容易忽略的点是“右图左移”不是“左图右移”。编码时最好统一方向否则后续可视化会经常错位。实际代码里常用torch.cat((left_feat, shifted_right_feat), dim1)循环D次实现。PSMNet用的是concat式的代价体简单直接但有一个明显弱点对于光照差异较大的场景concat特征对噪音的鲁棒性不够好。后续GwcNet把concat改成了分组相关groupwise correlation把左右特征分成若干组做点积效果更好、显存也更友好。不过concat的优点是信息保留完整网络自己学会怎么比较不必依赖预设的相似度度量所以在训练数据足够强的情况下PSMNet的基础能力依然很扎实。3.2 3D沙漏的编码-解码循环逻辑代价体有了接下来要在D、H、W三个空间维度上做全局聚合。这里不用2D卷积因为2D卷积只能在图像平面传播信息视差维度的邻居关系被完全割裂了。视差维度的“邻域”本身有强先验——真实场景中相邻像素的视差不会突变同一物体表面视差应该平滑变化。用3D卷积可以把空间平滑性和视差连续性一起建模这就是PSMNet把聚合网络做成3D结构的原因。具体结构上3D沙漏网络是一个编码器-解码器结构编码器这边做4次3D卷积下采样每次把空间分辨率减半、通道数翻倍中间再经过几层3×3×3卷积提取高语义特征解码器这边用3D反卷积或转置卷积逐步恢复分辨率并且引入跳跃连接把下采样前的细节拼回去。整个过程像是一个“压缩-理解-重建”的循环。很多刚接触的人不理解为什么PSMNet要堆叠多个沙漏。堆叠的意义在于让代价体经过第一轮全局推理后再进入第二轮修正。第一轮会产生一个粗糙但大范围的视差估计第二轮在这个估计的基础上重点修正细节相当于两阶段推理。论文里堆叠了3个沙漏结构每一层都会输出一个视差预测图计算损失时三个输出同时受监督我在实践里发现中间层的权重可以调低一点不然前期训练容易动摇。沙漏网络输出之后网络会在视差维度上做softmax概率最高的视差值作为预测结果。准确说PSMNet是做了加权soft argmin用概率分布和视差水平轴做加权和而不是直接取argmax这样可以得到亚像素级别的视差精度。3.3 深度尺度与显存预算的平衡3D代价体聚合是PSMNet的精度来源也是最大的显存杀手。我按B4、输入尺寸256×512、D192来估算代价体本身是4×32×192×64×128×4字节约400MB第一层3D卷积直接把这个张量扩了通道峰值显存经常飙到10GB以上。所以单卡训练时很多人只能把batch降到1到2或者把输入裁剪成192×384否则会直接爆显存。my实测在一个12GB的旧卡上256×512输入、batch1勉强能塞下反向传播时经常在最后几个3D反卷积层OOM。解决方式有三种裁剪输入训练、用梯度累积模拟较大batch、或者使用混合精度训练。混合精度对3D卷积的加速尤其明显代价体构造成为了另一个瓶颈因为那个循环拼接操作没法完全走半精度所以实际提速没想象中夸张。显存预算不够时也可以考虑把D降到128大多数室内场景足够用但室外远距离物体会被截断精度掉0.4个百分点左右。这个取舍要在项目开始前就想清楚不要等到训练完了才发现。4. 训练PSMNet的实操一条龙从合成数据预训练到真实数据微调4.1 预处理里面隐藏的“数据泄露”风险PSMNet最常用的训练套路是先在大规模合成数据集上预训练再到真实场景数据集上微调。合成数据量大、标注像素级精准但和真实世界存在域差距。这个阶段预处理有大量细节其中最容易忽视的是数据预处理时空一致性。左右图翻转、色彩抖动、随机裁剪这些操作如果只对单图做就把左右图像的极线关系破坏了。立体匹配训练里的数据增强必须保证左右图经过完全相同的空间变换。比如随机水平翻转时左右图要一起翻转并且注意视差映射也要跟着翻转左图变右图视差d变成原来的W-1-d随机裁剪要从左右图的同一位置裁剪必要时还要调整裁剪区域的视差偏移。另一个坑叫标注口径不统一。合成数据里有些区域标注是0表示无效有些数据集里0反而表示贴合近景。如果用错了掩膜把无效区域的0当成有效微差来训练损失会引导网络预测出大量伪0边缘瞬时段全糊掉。正确的做法是训练阶段用一个“valid mask”把视差0或者超出下标的像素从损失计算中刨除。4.2 损失函数为什么非平滑L1不可PSMNet的损失是平滑L1损失(smooth L1 / Huber loss)在预测视差与真实视差的绝对值小于1时使用平方项大于1时使用线性项。为什么要这样做如果只用L2损失遇到异常值比如标注错误或遮挡边界时会产生大梯度训练不稳定只用L1损失亚像素级别的细小误差不容易收敛到光滑解。平滑L1是在二者之间做平衡小误差时二次收敛大误差时就当线性问题处理。多阶段监督时三个沙漏输出分别算损失再加权求和。我在实际项目里把第一、第二层沙漏的损失权重设为0.5和0.5最后一层设为1.0这样前期阶段负责“搭骨架”最后一层负责“抠细节”。如果前三层权重都一样有时会出现回传梯度过大、网络整体偏保守的现象。4.3 微调阶段的BN冻结与学习率调度预训练之后进入微调阶段有一件容易被忽略但影响很大的事BatchNorm的统计量到底能不能继续更新。合成数据和真实数据的分布差距很大如果BN的running mean和running var继续随真实数据更新前几个batch很可能把统计量拉偏造成精度抖动。我的建议是微调的前几千步保持BN统计量的更新等网络适应真实数据分布后再手动冻结BN层也就是把affine参数固定、track_running_stats设为False继续训练主干和聚合网络。当然这个做法在不同实现里有争议我自己的多次实验显示冻结后测试精度更稳尤其当真实数据集规模很小几百张图时差异非常明显。学习率调度方面微调阶段建议用Adam优化器初始学习率从1e-3降到1e-4再用余弦退火或阶梯衰减。之前在合成数据预训练阶段用的是1e-3加80个epoch的余弦衰减迁移到真实数据后其实不需要从头来过15000步基本就能收敛。这里要避开的坑是微调时把整个网络的所有层都解开前面特征提取层很容易过拟合到真实数据的小样本上导致泛化变差。某些项目里会冻结特征提取器只训练3D聚合部分效果反而更好。5. 复现踩坑实录那些文档里不会写的“黑话”5.1 显存爆掉的现场处理跑3D代价体聚合网络大家最心烦的就是显存不够。我当时用一个12GB的卡加载预训练权重后准备微调一个batch切到2就直接OOM。试过减少D、裁剪输入、开混合精度、梯度累积几种方法最后将batch设为1、输入改成320×640、采用梯度累积4步来模拟batch4的效果总算跑起来。这里有一个技巧把代价体的构建部分和聚合网络分开计算显存。具体做法是先循环构建完所有视差切片存到显存再一次性送入3D网络这样虽然中间占用没减少但至少态势清楚。另一个更有效的工作是使用torch.utils.checkpoint在3D沙漏的编码器部分做激活检查点可以省掉一大半的中间激活显存代价是每次反向传播会重新做一遍前向计算训练时间总体增加10%左右。如果只是微调梯度累积加checkpoint是最稳的组合。5.2 预测图边缘出现“拖影”的排查复现PSMNet之后最容易观察到的问题物体边缘和遮挡区出现半透明的“拖影”就像视差预测在前后景之间做了连续插值。这是softmax加权和机制带来的必然结果。在遮挡区域或深度剧变处代价体的概率分布往往只有一个宽而平的峰soft argmin输出就会落在前后景之间的中间位置。排查时不要一上来就调网络结构。先检查输入图像有没有这一步——左右图有没有做归一化到[0,1]而不是[-1,1]归一化方式不同会导致特征分布偏离预训练权重概率峰变平。其次观察代价体在这一区域的概率分布形态如果确实是宽峰可以尝试把softmax的温度参数调低一点让分布更尖锐。最后再考虑后处理中值滤波或双边滤波对边缘拖影有一定抑制作用但只能治标。如果拖影集中在小物体周围改善空间最大的是数据。弱监督或半监督训练数据里小物体占的比例低网络没有足够的样本学会“小物体边缘处视差应该突变”这时可以在数据增强时加一些随机透视变换、局部放大增强小尺寸物体的出现频率。5.3 用错误掩膜评估精度——你以为的1.2%可能没那么美评估精度的环节也暗藏玄机。公开数据集的评估指标通常分前景、背景、全部三类分别计算误差像素占比。如果你照抄评估代码但掩膜生成逻辑不对比如没把裁剪造成的边界区域排除掉模型在实际应用中会露出原形。我自己就栽过一次把某公开无人驾驶影像上微调后的模型放到自己的室内数据集上测试测出来D1误差只有1.8%感觉不错但输出的视差图一看近处桌椅边界明显糊。查了半天才发现评估时把近处截断区域视差大于D_max也算进去了这些像素因为视差数值超标自动失败拉低了指标。正确做法是构建有效掩膜时把真值视差超出预测范围、位于图像边界裁剪带内的像素全部排除再分层统计。除此之外还要检查模型是否在“作弊”因为网络输出的是argmax加权和如果标签分辨率低、真值本身在某些区域有量化误差模型会学到略微偏向中心值的输出。这类偏差只能通过修改损失或者增加亚像素监督来治。6. 从PSMNet往前再走一步后续方案怎么选型6.1 变体对比哪些坑被填上了PSMNet既经典又昂贵后续工作大多在它基础上改了某个核心模块。做一个量化对比表会更直观一些。方法核心改动精度/速度特点显存需求适用场景建议PSMNet原始SPP3D沙漏精度高速度慢极高速度不敏感、需要高精度场景GwcNet分组相关代价体精度略优显存更低高蒸馏或精度要求很高的离线任务AANet可变形聚合替代3D卷积速度大幅提升精度略降较低实时或近实时应用ACVNet注意力机制过滤代价体精度优对遮挡更好高复杂遮挡、动态场景RAFT-Stereo迭代GRU更新视差精度优显存比3D代价体低中有迭代时间预算追求端到端效果IPDAETransformer全局建模对弱纹理和遮挡最好高只看重精度不在乎硬件成本以我的实际感受来说如果目标是给后续的3D重建或者自动驾驶感知模块提供深度PSMNet仍然是一个很稳的基座因为它的错误分布比较“温和”不会出现大片黑色空洞。而AANet这类轻量方法更适合Demo环境放到真实产品里细节不够。6.2 什么时候继续用PSMNet我会建议在以下三种情况继续选择PSMNet第一你已有预训练权重需要在自有小数据集上微调并快速验证效果第二你的输入分辨率不高小于512宽显存可以接受希望拿到当前精度最扎实的baseline第三你的项目里需要和后续的GwcNet、ACVNet等做精度对比实验PSMNet作为经典对照很可信。如果目标场景是无人机实时深度、嵌入式部署或者视频流的连续估计那就不应该停在PSMNet上而应该去研究轻量化的代价体聚合方案或者直接上迭代式更新网络。这里没有绝对的好坏只有取舍。如果想在PSMNet基础上做增量改进我最推荐的两个方向一是把代价体从concat换成groupwise correlation改动不算大但能明显缓解光照差异带来的误匹配二是把三个沙漏结构改成两个在显存和训练时间上都省不少精度损失可控制在0.2个百分点以内。在我自己的项目里最后保留了PSMNet作为离线高精度方案的基线同时搭了一个轻量网络做实时预览。这种双轨思路在很多实际工程里都适用先用PSMNet预标注一批数据再用轻量网络做在线推理既能保证质量又能满足速度要求。如果你正卡在“精度够但速度不够”或者“显存不够但效果要再提”这种两难问题上我建议不妨先跑通PSMNet这一条完整链路——把特征、代价体、聚合、监督方式都吃透了后面换任何结构都会顺手很多。