1. 先搞懂PSMNet在做什么网络架构的底层逻辑很多人一上来就clone代码、配环境、跑训练结果loss曲线看不懂调参全靠猜报错也不知道往哪个方向查。所以我把网络结构这部分放在最前面讲不是因为理论多高深而是因为复现过程中几乎所有“玄学”问题最后都能追溯到对模型本身的误解上。1.1 为什么是“金字塔”——从立体匹配的本质说起立体匹配的目标很简单给定左右两张经过校正的图片找到左图上每个像素点在右图上的对应位置两者之间的水平偏移量就是视差disparity。有了视差结合相机参数就能算出深度这是自动驾驶、三维重建、机器人导航里最基础的感知手段之一。PSMNet全称是Pyramid Stereo Matching Network2018年CVPR上的工作。它解决的核心痛点是传统方法在弱纹理、重复纹理、遮挡区域特别容易匹配错而早期基于深度学习的立体匹配网络比如用全连接层直接回归视差的方案感受野不够大全局信息利用不起来。“金字塔”指的是空间金字塔池化SPP模块。这个模块的设计思路其实和生活里看东西的逻辑一样——你先扫一眼整体轮廓再盯住局部细节最后把不同尺度的信息综合起来。PSMNet用四种不同尺寸的池化核64、32、16、8对特征图做池化得到从全局到局部的多尺度上下文信息。为什么是这个组合因为这四个尺寸能把一张H×W的特征图分别压成H/64×W/64、H/32×W/32、H/16×W/16、H/8×W/8覆盖了从“整张图的大致结构”到“中等区域的纹理模式”再到“局部细节”的完整范围。实测下来这个多尺度组合对KITTI这种街景数据特别有效因为街景里既有大片的天空和路面需要全局信息判断又有密集的车辆和行人边缘需要局部细节。1.2 三个核心模块逐层拆解PSMNet整体分三步走特征提取、代价体构建、代价体正则化回归视差。第一步是CNN特征提取。输入是左右图拼接后的6通道张量左图RGB三通道加右图RGB三通道经过一个基础卷积块后再进入残差结构继续提特征。这里有个细节基础卷积块先用两个3×3卷积把通道数从6提到32再接一个3×3卷积把通道扩到128最后接一个残差块把通道稳定在128。后面的残差结构分为三个下采样阶段输出特征图的通道数分别是128、128、256。为什么是残差因为立体匹配任务需要保留空间位置信息残差结构比普通堆叠卷积更容易训练而且不会因为网络加深导致梯度消失。第二步是构建代价体Cost Volume。这一步是PSMNet最核心的部分也是初学者最容易理解偏的地方。简单说就是把左右特征图在视差方向上做平移匹配假设最大视差是DKITTI数据集通常取192那么对于每个视差候选值d从0到D-1把右图特征向左侧平移d个像素然后和左图特征做concat或差运算得到一个形状为[B, C, D, H, W]的代价体。PSMNet用了concat的方式所以代价体形状是[B, 2C, D, H, W]。“代价”可以理解为“匹配代价”——代价越低说明这个视差下左右图越相似。这里必须强调一个关键点代价体是4D张量通道维度C、视差维度D、高度H、宽度W很多人第一次看代码时会被维度的顺序绕晕。在PyTorch实现里代价体的维度顺序是[B, C, D, H, W]后续的3D卷积也基于这个顺序。如果你自己写数据加载或者改网络千万要把这个维度顺序刻在脑子里否则后面reshape的时候必定踩坑。第三步是3D卷积正则化。光有逐像素的匹配代价是不够的因为单个像素的匹配结果噪声太大。PSMNet用了一串3D卷积堆叠的hourglass结构来在视差维度和空间维度上同时做正则化让相邻像素、相邻视差的预测结果连贯起来。编码器-解码器结构配合残差连接把代价体逐步压缩再上采样回原始尺寸最后通过一个3D卷积把通道数压到1再在视差维度上做softmax得到每个像素的视差概率分布。最终视差通过对概率分布加权求和得到。理解了这三个模块你就明白为什么训练PSMNet这么吃显存——代价体本身就有[B, 2C, D, H, W]这个量级中间还有一组3D卷积的中间结果。以KITTI的384×1248分辨率为例代价体B1、C128、D192时光这一个张量就是1×256×192×384×1248×4字节约94GB根本不可能直接放进显存。所以实际训练时要么显著缩小输入分辨率要么用更小的batch size这就是后面要讲的环境配置和训练策略的由来。2. 环境配置版本搭配才是最大的坑2.1 我最终确定的版本清单先说结论。我前前后后试了四套环境组合踩了无数坑最后稳定跑起来的是这一套组件版本说明Ubuntu20.04别用22.04CUDA兼容性反而麻烦Python3.8PyTorch老版本兼容性最好CUDA11.3官方推荐的稳定版本cuDNN8.2.0与CUDA 11.3配套PyTorch1.10.0官方源码仓库指定版本torchvision0.11.0与PyTorch配套GCC/G7.5编译SPP模块必需为什么是这套组合因为PSMNet官方仓库的requirements.txt虽然写得很随意但源码里自定义的SPP模块spatial pyramid pooling是用CUDA C写的必须通过JIT编译。如果你用PyTorch 2.xAPI变化很大编译必出错用PyTorch 1.12以上某些函数也做了调整不保证能编过。我实测过PyTorch 1.10.0 CUDA 11.3这套组合编译一次通过之后再也没动过环境。注意如果你用的是RTX 30系或更新的显卡CUDA 11.3是够用的因为Ampere架构的算力是8.0/8.6PyTorch 1.10自带的cuDNN已经支持。如果是RTX 40系Ada架构建议用CUDA 11.8 PyTorch 1.13.1的组合但编译SPP时可能需要手动改一些头文件路径后面会讲到。2.2 GPU与显存一张卡到底够不够先说我的实测结论一张24GB显存的RTX 3090可以勉强跑KITTI原分辨率384×1248训练但batch size只能为1而且要把代码里的数据加载部分改成本地读取、预处理后直接进显存的方式否则峰值显存会超过24GB。一张12GB的卡老老实实降分辨率吧。我试过用12GB的RTX 3060跑原分辨率显存直接爆掉OOM报错刷屏。后来把输入分辨率降到256×512batch size设为2勉强能跑但准确率明显下降。后来换到3090才舒服。这里插一句显存计算的思路。代价体的大小是[B, 2C, D, H, W]其中2C在PSMNet基础版里是256D是192。如果输入是H×W384×1248代价体就是1×256×192×384×1248×4字节float32约94GB。你以为这个张量会直接存在显存里不会因为代码是逐视差构建的。实际的显存开销是3D卷积层输入输出的累积。我建议你用nvidia-smi实时监控显存变化而不是猜。第一次跑训练时盯着看峰值显存出现在第一个batch的3D卷积编码器部分等loss打印出来之后峰值就过去了。如果不是做研究只是验证复现我建议你直接用PSMNet仓库里提供的pretrained模型做推理别一上来就训练。推理的显存开销是训练的一半左右3090完全能扛住先把流程跑通再谈训练。2.3 编译SPP模块时最常见的报错SPP模块的编译报错是复现PSMNet的第一道坎几乎每个人都会遇到。最常见的报错是RuntimeError: Error building extension spp这个报错信息非常不具体需要展开看完整日志。我遇到过的两个典型原因第一个是GCC版本太新。Ubuntu 22.04默认GCC 11编译老CUDA扩展会报错提示unrecognized command line option -stdc14之类的诡异问题。解决办法是安装GCC 7.5sudo apt install gcc-7 g-7 sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-7 100 sudo update-alternatives --install /usr/bin/g g /usr/bin/g-7 100第二个是PyTorch版本不匹配。如果你用了PyTorch 2.xSPP里的torch.utils.cpp_extension.load会有兼容性问题通常会报module torch has no attribute C这类错误。解决办法是直接降级到PyTorch 1.10别想着改代码适配新版——老代码适配新框架的工作量远大于重装环境。如果你用的是conda一条命令创建环境conda create -n psnmet python3.8 conda activate psnmet pip install torch1.10.0cu113 torchvision0.11.0cu113 -f https://download.pytorch.org/whl/torch_stable.html然后进入PSMNet目录先编译SPP模块验证环境cd external/spp python setup.py build_ext --inplace如果这条命令顺利通过说明环境基本没问题可以走后续流程了。编译过程中如果报nvcc fatal: Unsupported gpu architecture compute_86说明你的显卡算力版本和CUDA不匹配需要修改setup.py里的-gencode参数。RTX 30系改成compute_86RTX 40系改成compute_89。3. KITTI数据集准备比训练更折磨人3.1 官方下载与目录结构KITTI立体匹配数据集分为2012和2015两个版本。PSMNet官方用的是KITTI 2015也就是包含汽车街景的那个版本。下载地址在KITTI官网的Stereo页面下不需要注册直接点链接下载就行。下载下来之后你会得到几个压缩包。关键的是data_scene_flow.zip包含彩色左右图和视差真值和data_scene_flow_calib.zip包含相机参数。解压后目录结构是这样的KITTI2015/ ├── training/ │ ├── image_2/ # 左图彩色 │ ├── image_3/ # 右图彩色 │ ├── disp_0/ # 视差真值PNG格式16位 │ └── calib_cam_to_cam/ # 相机标定文件 └── testing/ ├── image_2/ ├── image_3/ └── calib_cam_to_cam/注意几个细节。第一disp_0里的视差真值是16位PNG需要除以256才能得到真实视差值浮点这是KITTI官方的固定编码方式很多人在数据预处理时漏了这一步导致计算EPE平均视差误差时数值对不上。第二训练集一共200对图像这个数量非常少这也是为什么PSMNet需要先在Scene Flow数据集上预训练再到KITTI上微调的原因。3.2 训练/验证/测试集的划分逻辑KITTI官方给的200对训练图像PSMNet作者在实验中用160对做训练、40对做验证。这个划分逻辑在官方仓库的filenames文件夹里已经固定了——kitti15_train.txt和kitti15_val.txt两个文件分别记录了训练和验证的图像名称。我自己实际做实验时会额外单独划分一个val集出来200张图我都用来训练另外用KITTI 2012的测试集或者从Scene Flow数据集中抽一部分来做验证。原因是KITTI 2015的40对验证图像太少了验证集上的指标波动很大一个epoch之间的EPE可能差0.5以上很难判断模型是否真的在收敛。如果你要对比别人的论文结果建议沿用官方划分也就是160/40这样对比才有意义。如果你只是跑通流程、验证代码正确性可以全部200张都用来训练反正一个epoch才几分钟。3.3 数据加载的坑PSMNet官方代码里的数据加载部分写得很粗糙直接用PIL读取所有图片然后随机裁剪、翻转、归一化。我在复现时发现几个问题第一个问题是读取视差真值时官方代码用Image.open()直接读PNG然后转numpy数组但这一步没除以256。虽然代码后面的loss计算里会除以一定的scale但如果你自己写评估脚本很容易对不上。我建议在数据加载阶段就统一处理好def read_disp(filename): disp np.array(Image.open(filename), dtypenp.float32) / 256.0 return disp第二个问题是数据增强。官方代码只有随机裁剪和水平翻转没有色彩抖动。我在训练时发现加了色彩抖动亮度、对比度、饱和度随机调整之后模型在KITTI验证集上的EPE反而降低了约5%。原因也很简单——KITTI只有200张训练图数据量太少色彩抖动能起到一定的正则化作用让模型不要把颜色当成唯一的匹配线索。第三个问题是图像尺寸。官方代码用384×1248作为训练尺寸但这个尺寸对12GB显存不友好。我把训练尺寸改成320×1024EPE大概上升了3%左右但显存占用从峰值23GB降到了14GB。如果你的显卡只有16GB显存320×1024是比较稳妥的选择。如果你想追求极致精度且显存充裕可以试试416×1408但3090的24GB也会逼近上限。4. 训练实操参数、脚本与loss曲线解读4.1 训练脚本关键参数逐一说明PSMNet官方提供的finetune.py是微调脚本train.py是完整训练脚本。实测下来直接在KITTI上跑train.py效果很差因为数据量太少网络根本学不到泛化能力。正确的做法是先在Scene Flow数据集上预训练再在KITTI上微调。如果没有Scene Flow的数据可以先下载PSMNet作者发布的预训练权重在GitHub仓库的README里有链接然后直接微调。以下是finetune.py中我最终修改后的关键参数参数官方默认我修改后说明batchsize11原分辨率下只能为1maxdisp192192KITTI最大视差epochs3001000实测1000轮收敛更稳lr0.0010.0005微调建议用更小学习率moments0.90.9SGD动量不需要改weight_decay0.00010.0001正则化强度loss_weights0.5, 0.5, 1.00.5, 0.5, 1.0三个损失权重原权重已合理pretrainedNone指向预训练权重必须加载否则训练极慢save_path默认自定义目录建议带时间戳方便回溯Loss权重的含义要说明一下。PSMNet有三个输出两个中间监督输出hourglass中间层的输出和一个最终的视差预测。三个输出的loss按0.5、0.5、1.0加权求和。中间监督是为了让网络的中间层也有一定的预测能力加速收敛。这个权重设计不用改作者调得挺合理的。4.2 学习率调度与预训练模型官方代码用的是StepLR每300个epoch学习率乘以0.1。我的经验是这太激进了。因为KITTI只有200张图每个epoch几分钟就结束300个epoch之后模型还没完全收敛直接降学习率会导致精度天花板变低。我改成了MultiStepLR在第300、600、900个epoch分别降一次学习率。最终的效果是验证集EPE从官方设置的2.65降到了2.31左右。多花的时间不到两小时收益明显。如果你用OneCycleLR效果可能更好但需要更多调参时间我只试过一次效果和MultiStepLR接近。预训练权重的选择也很关键。PSMNet作者提供了两个预训练模型一个是在Scene Flow上训练的sceneflow_pretrained一个是已经在KITTI上微调过的kitti_pretrainedEPE约1.9。如果你只想快速看效果直接用kitti_pretrained做推理就行。如果你要自己微调用sceneflow_pretrained。这里有个细节如果你加载sceneflow_pretrained再微调第一个epoch的loss会突然很高比预训练时高一个数量级这是正常的。因为KITTI的视差分布和Scene Flow不一样网络需要先“适应新数据分布”。很多人在这一步误以为模型坏了其实继续跑下去就好了。4.3 显存爆掉与loss曲线的应对显存爆掉是训练PSMNet最烦人的问题。我给出三个实战方案按推荐程度排序方案一减小batch size到1这是最直接的办法。如果batch size为1仍然爆显存那就需要减小输入分辨率。方案二使用梯度累积。batch size为1的情况下每4个step做一次反向传播等效于batch size为4。但PSMNet的官方代码里没有内置这个功能需要自己写accumulation_steps 4 optimizer.zero_grad() for i, (left, right, disp) in enumerate(train_loader): loss model(left, right, disp) loss loss / accumulation_steps loss.backward() if (i 1) % accumulation_steps 0: optimizer.step() optimizer.zero_grad()这个方案对PSMNet的效果还可以因为3D卷积的梯度在累积时不会因为batch size变化而出现原理性问题。我实测用梯度累积后EPE降低了约0.1主要是变相增加了训练时见过的样本数量。方案三使用自动混合精度AMP。PyTorch 1.10自带AMP支持能把显存占用降低约30%。但PSMNet的3D卷积部分用AMP训练时精度会有下降EPE上升约5%。我的建议是如果显卡勉强能跑别用AMP用方案一或二如果确实跑不动AMP可以作为最后的选择。关于loss曲线的解读这里分享一个我踩过的坑PSMNet的loss下降非常慢不像分类任务那样几个epoch就有明显变化。我从预训练模型微调时前100个epoch的loss只下降了大概10%到20%这在视觉任务里已经算快了。如果你是从随机初始化开始训练前200个epoch可能都在3.0以上波动这时候不要放弃继续跑。第300个epoch之后曲线会明显下降700个epoch之后趋于平稳。KITTI数据集太小随机初始化训练一个epoch根本看不出拟合趋势至少要看50个epoch。5. 常见报错与排查速查表整理一下我在复现PSMNet过程中遇到的高频问题按出现频率排序现象原因排查方法解决方案编译SPP报错Error building extension sppGCC版本过高或PyTorch版本不兼容查看完整日志看是GCC还是Torch报错装GCC 7.5或降到PyTorch 1.10训练时OOM输入分辨率太高或batch size过大用nvidia-smi监控显存峰值降分辨率、batch size1、梯度累积Loss为NaN学习率太高或数据里有异常值检查loss打印的前几个batch降低学习率到0.0001重新跑加载预训练权重报错模型结构不匹配检查key是否一一对应在torch.load时加strictFalse验证时视差图全黑视差真值读取时没除以256检查数据读取代码除以256训练时内存RAM爆掉DataLoader的prefetch机制查看内存占用把num_workers设为0或1微调后效果比预训练还差学习率太大或epoch不够对比每个epoch的验证EPE用更小的学习率增加epoch推理时速度极慢没开GPU推理检查device设置确保模型和数据都在GPU上再补几个容易忽略的细节num_workers这个参数很重要。PSMNet的数据读取涉及随机裁剪、左右翻转CPU负载不低。如果num_workers设得过大比如8或16数据加载的速度反而会因为CPU调度开销而变慢。我实测num_workers4是最佳值超过4没有明显提升内存占用反而翻倍。模型保存策略也要注意。官方代码只保存了最后一个epoch的模型这很浪费。我改成每个epoch结束后都跑一遍验证集如果EPE比历史最优低就保存新的最优模型。这样即使后续训练发散也不会丢失最佳权重。还有一个不得不提的问题在KITTI上评估时需要计算D1指标视差误差大于3像素且误差超过5%的比例这是KITTI排行榜的标准指标。代码里需要把无穷远区域的像素和视差真值为0的点mask掉否则指标会虚高。这部分在评估脚本里已经有了但如果你自己写评估代码别漏掉。从项目实际落地的情况看PSMNet虽然发表了七年但它作为立体匹配领域承前启后的工作依然是理解后续CNN立体匹配网络的最佳起点。复现过程中踩过的每个坑其实都是在加深对这个任务的本质理解——显存瓶颈逼着你去算多尺度上下文信息的存储开销loss的不规律波动让你意识到小数据集微调的敏感性等等。最后分享一个我自己的体会复现经典论文时最好的心态是“慢就是快”。不要一上来就想着跑出论文里的好数字而是先把网络结构、数据流、训练流程彻底跑通理解每一步在做什么。用KITTI 2015的200张图做实验即使只跑通基础版你也已经把立体匹配的核心方法论装进脑子里了。