PSMNet复现全流程踩坑记录：从环境配置到KITTI训练

搞立体匹配方向的同学，一定绕不开PSMNet这个经典模型。我在两周内从零复现它，期间踩了环境配置、数据读取、显存爆炸、指标计算等各种坑，最后总算把它在KITTI上完整跑通。这篇文章把我整个踩坑过程，从环境搭建到数据集训练，全部梳理成一份可直接上手的操作记录，里面包含了我实际用的命令、代码片段和排查思路。无论你是刚入门深度学习想找个经典模型练手，还是已经在做立体匹配相关课题，这份记录应该都能帮你少走不少弯路。

1. 先看一遍PSMNet再动手：核心思路与复现价值

1.1 为什么要复现这个2018年的老模型

虽然PSMNet是2018年的工作，但直到今天，KITTI立体匹配排行榜上依然能看到它的身影，很多后续方法也都拿它当baseline对比。对新手来说，复现PSMNet能一次性接触到立体匹配里最核心的几个套路：特征提取、代价体构建、3D卷积聚合、soft argmin回归视差。这套pipeline搞明白了，再看RAFT-Stereo、IGEV这类新模型会轻松很多。

我自己复盘整个过程，最大的感受是：复现经典论文的价值不在于把指标刷多高，而是把“读论文”变成“跑通代码再回头看论文”，很多抽象概念在亲手实现之后才有实感。比如cost volume这东西，看论文时觉得不就是特征拼接嘛，真正自己去写那个循环shift左特征和右特征的代码时，才明白为什么它是整个模型里最吃显存的部分。

1.2 PSMNet的核心结构拆解

PSMNet的大致流程是这样的：左右两幅图先经过一个共享权重的CNN提取特征，输出分辨率是原图的四分之一，通道数一般是32。然后特征进入空间金字塔池化SPP模块，用不同尺寸的平均池化获取多尺度上下文信息。接下来是关键的cost volume构建：对于每个视差值d，把左特征和右特征在宽度方向上做平移后拼接，得到形状为[B, 2C, D, H/4, W/4]的代价体。之后用堆叠的3D卷积对这堆代价体做正则化，最后在视差维度上做softmax，加权求和得到亚像素精度的视差图。

这里有几个复现时必须注意的细节。SPP模块的输出尺寸对齐是个隐藏坑，不同池化核大小算出来的特征图尺寸可能差一个像素，直接上采样回原尺寸再接1x1卷积容易出问题。我的建议是，把池化层的stride和kernel size都设置成能整除的数值，或者干脆用AdaptiveAvgPool2d，省心很多。

另外就是maxdisp这个超参数。KITTI数据集的视差范围通常不超过192像素，所以大多数实现默认maxdisp=192。这个值决定了cost volume在视差维度上的尺寸，直接影响显存占用。复现时先按192跑，万一你的应用场景视差范围特别大，再调整也不迟。

1.3 三个“推理正确但跑不通”的设计细节

第一个是视差回归。PSMNet不是直接输出一个标量视差，而是对cost volume在视差维度上做softmax得到概率分布，然后对每个视差级别做加权求和。这个操作在论文里叫soft argmin。实现起来很简单，几行代码就能搞定，但有个容易出错的地方：softmax的logits如果不做缩放，概率分布可能会过于尖锐或平坦，导致训练初期梯度异常。有些实现会在softmax之前除以一个温度系数，实测下来更稳。

第二个是中间监督。原版PSMNet的3D CNN部分包含堆叠的hourglass结构，会在不同阶段输出视差估计并加上损失。复现时如果仓库代码里有这种多尺度监督，别急着删掉，它对收敛速度有明显帮助。我第一次复现时把中间监督全去掉了，结果训练到200个epoch时loss还降不到理想值，加回去之后明显好转。

第三个是输出分辨率的问题。网络最后输出的视差图是原始分辨率的四分之一，计算损失前需要上采样回原图大小和ground truth对齐。这个上采样一般用双线性插值就能搞定，但要注意：如果你后面要做左右一致性检查或者后处理，最好在四分之一分辨率上先做，再上采样，不然边缘会糊。

2. 环境配置的深坑：Python、PyTorch、CUDA与扩展模块

2.1 版本组合选择：不要装最新，要装稳定

环境配置是第一个让我抓狂的环节。PSMNet这个项目本身不算复杂，但它依赖的PyTorch版本和你显卡驱动之间必须匹配，否则会出现各种莫名奇妙的报错。我最终选定的组合是：Python 3.8 + PyTorch 1.13.1 + CUDA 11.7，跑得非常稳。

在动手装环境之前，第一步一定是先看自己显卡驱动支持的CUDA版本，用nvidia-smi查看右上角的CUDA Version。注意，这里显示的版本是驱动支持的最高CUDA版本，不代表你机器上已经装了对应CUDA工具包。PyTorch安装时选择CUDA 11.7还是11.8，只要不超过驱动支持的最高版本就行。

如果你用的是Anaconda，创建环境的命令大概是这样的：

conda create -n psmnet python=3.8 -y conda activate psmnet pip install torch==1.13.1+cu117 torchvision==0.14.1+cu117 --extra-index-url https://download.pytorch.org/whl/cu117

装完之后，用一行命令验证GPU是否可用：

python -c "import torch; print(torch.__version__, torch.cuda.is_available(), torch.cuda.get_device_name(0))"

我当时遇到过一种情况：torch.cuda.is_available()返回True，但真跑模型时报CUDA error: no kernel image is available for execution on the device。这通常是PyTorch版本和显卡架构不匹配导致的，比如太老的显卡装了太新的PyTorch。遇到这种问题，就只能换一个兼容你显卡架构的PyTorch版本，没有别的捷径。

2.2 CUDA扩展编译：gcc和nvcc的版本战争

市面上很多PSMNet的开源实现，除了纯PyTorch代码之外，还需要编译一些自定义的CUDA扩展，比如resample2d、channelnorm之类的模块。这些扩展在训练和推理时用于处理特征图对齐，功能上用纯PyTorch也能实现，但原版代码为了追求效率选择了CUDA实现。

编译这类扩展，最坑的就是gcc版本和nvcc不兼容。我自己遇到过一个经典报错：

error: #error -- unsupported GNU version! gcc versions later than 8 are not supported!

这个报错的意思是，你机器上的gcc版本太新，超出了当前CUDA版本支持的范围。解决办法有几种：

• 如果系统里装了多个gcc版本，可以用环境变量指定一个旧版本：export CUDAHOSTCXX=/usr/bin/g++-8
• 如果是conda环境，可以装一个低版本的gcc：conda install gxx_linux-64=7
• 如果源码是用setup.py安装的，也可以临时改一下编译参数

另一个常见报错是fatal error: cuda_runtime.h: No such file or directory，这是因为编译时找不到CUDA头文件。解决办法是设置CUDA_HOME环境变量：

export CUDA_HOME=/usr/local/cuda export PATH=$CUDA_HOME/bin:$PATH export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH

2.3 其余依赖包清单

除了PyTorch之外，还需要几个常用库：opencv用于图像读取，tensorboardX或tensorboard用于可视化训练曲线，matplotlib用于画图，tqdm用于显示进度条。一次性装齐可以省很多事：

pip install opencv-python tensorboardX matplotlib tqdm

这里有个小提醒：opencv-python和opencv-contrib-python不要同时装，会冲突。读KITTI的16位视差图用普通版opencv就够了。

IDE方面，我用的VSCode配合Remote SSH，直接在服务器上改代码、看日志、可视化都挺方便。如果你也是用VSCode，记得装Python扩展和Pylance，调试断点功能对排查数据加载问题很有帮助。

3. KITTI数据集的下载、读取与加载器实现

3.1 数据集的目录结构和训练/测试划分

KITTI立体匹配数据集分2012和2015两版，我复现时用的是KITTI 2015。下载完解压之后，目录结构是这样的：

data_scene_flow/ training/ image_2/ 左目彩色图像 image_3/ 右目彩色图像 disp_occ_0/ 带遮挡标注的视差图 disp_noc_0/ 非遮挡区域的视差图 testing/ image_2/ image_3/

training目录下共有200对图像，testing目录下也有200对左右。图像分辨率大约是1242x375，整体偏宽幅。没有官方指定的训练集和验证集划分，常见做法是把training里面的图像随机分成160对训练、40对验证，或者直接用一个固定随机种子划分，方便对比实验结果。

KITTI 2012的目录结构略有不同，用的是colored_0和colored_1来存放左右图，但处理逻辑基本一致。如果你需要对照排行榜上的历史结果，2012和2015都有完整的排名，可以挑一个基准数据来验证自己的复现效果。

3.2 视差图的正确读取方式

这是整个复现过程里最容易忽略但影响最大的一个细节。KITTI的视差图是16位PNG格式，真实视差值需要除以256才能得到浮点数。很多人第一次读的时候直接用cv2.imread默认参数，读出来是8位图，所有视差值都变成了0或者255，训练loss直接崩掉或者变成奇怪的值。

正确读取代码如下：

import cv2 import numpy as np def read_disp(path): disp = cv2.imread(path, cv2.IMREAD_UNCHANGED) disp = disp.astype(np.float32) / 256.0 return disp

disp_occ_0和disp_noc_0的区别在于，前者包含遮挡区域的标注，后者只标注非遮挡区域，遮挡区域像素值为0。训练时一般用disp_occ_0，配合一个mask把视差值等于0的无效像素过滤掉。评估时如果想和排行榜对齐，需要明确自己用的是非遮挡还是全部区域指标，这两个数值差不少。

我还踩过一个隐蔽的坑：有的读取代码用PIL.Image.open读取16位PNG，但PIL默认可能把它当成I;16模式，直接.astype(np.float32)不做除以256，导致loss异常大。所以不管用什么库读，一定记得除以256，并且做一次可视化确认。

3.3 Dataset与DataLoader实现要点

Dataset类的核心逻辑其实很简单：根据索引拿到对应的左图、右图、视差图，做一些预处理，最后返回三个张量。我实现时大概是这样的结构：

class KITTIDataset(Dataset): def __init__(self, root, file_list, transform=None): self.left_paths = file_list self.right_paths = [p.replace('image_2', 'image_3') for p in file_list] self.disp_paths = [p.replace('image_2', 'disp_occ_0').replace('.png', '_disp_occ_0.png') for p in file_list] self.transform = transform def __getitem__(self, idx): left = cv2.imread(self.left_paths[idx]) right = cv2.imread(self.right_paths[idx]) disp = read_disp(self.disp_paths[idx]) # 这里根据文件名映射需要仔细核对，KITTI的文件名后缀有点绕 # 比如 000000_10.png 对应的视差图是 000000_10_disp_occ_0.png if self.transform: left, right, disp = self.transform(left, right, disp) return left, right, disp

文件名的映射是另一个容易踩坑的地方。KITTI 2015的左右图和视差图文件名不是简单的前缀替换，而是类似000000_10.png对应000000_10_disp_occ_0.png这种多了一段后缀。我一开始想当然地用replace('.png', '.png')，结果一堆文件找不到，排查半天才发现是命名规则的问题。

一个更省事的方式是直接遍历目录，把文件名和对应的左右图、视差图路径存成列表。因为KITTI的数据集不大，没必要在__getitem__里做复杂的路径拼接，启动时把这个映射关系一次性构建好就行。

DataLoader部分有几个效率相关的设置：

train_loader = DataLoader( train_dataset, batch_size=8, shuffle=True, num_workers=4, pin_memory=True, drop_last=True )

num_workers设置成4或者8可以显著提升数据加载速度，但也不是越大越好。我曾在服务器上把num_workers设成16，结果训练启动时直接段错误，后来降到8就好了。pin_memory=True在GPU训练时能减少数据传输时间，建议开启。如果你发现训练过程中GPU利用率经常掉到很低，多半是数据加载跟不上了，优先检查这两项。

3.4 数据增强中的几何变换坑

训练时做数据增强有利于提高泛化性，但立体匹配任务的数据增强比普通分类任务麻烦，因为左右图和视差图三者必须同步变换，而且有些变换会改变视差语义。

我踩过最深的一个坑是水平翻转。如果只是简单地把左右图和视差图同时水平翻转，其实破坏了左右图的对应关系。正确的做法有两种：

• 进行水平翻转时，将左右图交换，同时将视差图水平翻转（视差值本身不变）
• 或者干脆不做水平翻转，只做颜色抖动、亮度调整、高斯噪声这类不会改变几何关系的增强

我后来偷懒选了第二种，只在颜色和亮度上做扰动，效果还不错。如果你想做水平翻转，强烈建议先用一对样本做可视化，确认翻转后的左右图和视差图语义是对得上的，再把它加进训练流程。

垂直翻转倒是安全很多，因为视差只跟水平方向有关，垂直翻转对视差值没有影响，左右图同时垂直翻转即可。

随机裁剪也是常用增强手段，裁剪左右图和视差图的同一区域，视差值不变。常见的裁剪尺寸是256x512或320x736，裁剪太小会丢失上下文信息，太大又会让显存压力剧增。

4. 训练流程、超参与显存优化实操

4.1 损失函数与评估指标怎么算

PSMNet的损失函数用的是Smooth L1 Loss，也就是Huber Loss，对离群点没那么敏感。实现时要注意用mask过滤掉无效视差像素：

import torch.nn.functional as F def smooth_l1_loss(pred, gt, mask): loss = F.smooth_l1_loss(pred, gt, reduction='none') loss = loss[mask].mean() return loss

评估指标方面，KITTI最常用的两个指标是EPE和D1。EPE是平均端点误差，即预测视差和真实视差的平均绝对像素差。D1是误匹配率，统计误差大于3像素且相对误差超过5%的像素比例。计算D1的代码大概长这样：

def compute_metrics(pred, gt, mask): pred, gt = pred[mask], gt[mask] epe = (pred - gt).abs() err = ((epe > 3.0) & (epe / gt > 0.05)).float().mean() return epe.mean().item(), err.item() * 100

如果你在验证集上看到D1在3%到5%之间，EPE在1像素左右，基本说明复现是成功的。原论文报告的结果会比随机验证集划分好一些，因为官方评测集和训练集分布更一致，不用强行追求完全一样的数字。

4.2 超参数设置背后的逻辑

我训练时用的关键超参数如下：

学习率策略上，原论文是每200个epoch衰减10倍。我在实际训练中发现，400个epoch的进度下，前200个epoch用0.001，之后切到0.0001，模型能比较好地收敛。这个策略不是最优的，但省心稳定，不至于在中途因为学习率太大而loss发散。

优化器我用的Adam，比SGD收敛快不少。对复现任务来说，Adam是最省事的选择。如果你的显卡显存紧张，batch_size被迫设得很小，可以考虑用梯度累积来模拟较大的batch size，效果更接近原论文的设定。

4.3 显存爆炸的处理手段

PSMNet最大的特点就是吃显存。我粗略算过：裁剪尺寸256x512，特征图四分之一分辨率就是64x128，通道数32，maxdisp为192时，cost volume本身的形状是[B, 64, 192, 64, 128]。光这一个张量，batch size为8时就要占用约100GB？不对，我重算一下。

每对特征的cost volume不考虑batch大小，是2倍通道数（左右特征拼接，所以是64通道），乘以192个视差级别，再乘以64x128的空间分辨率。所以单个样本的cost volume约是64x192x64x128 = 1亿个float元素，换算成内存大概是400MB左右。这是构建出来的代价体，还没算3D卷积的中间结果。batch size为8时，光这层就到3.2GB了，再加上后续3D卷积的特征图，24GB显存也就勉勉强强。

所以我复现时做了一个取舍：batch size从论文常用的16降到8，裁剪尺寸从320x736降到256x512。如果你只有12GB显存，建议batch size设为4，裁剪尺寸改成192x384，或者把maxdisp降到160先验证流程。

另外有个小技巧：不需要高频执行torch.cuda.empty_cache()，这命令看似释放了显存，实际上会降低性能。除非你在同一个进程里反复加载和释放大模型，否则让PyTorch自己管理就好。

4.4 训练日志与模型保存

训练时建议定期打印loss和验证指标，同时用tensorboard记录训练曲线。我实际跑出来的日志大概长这样：

epoch 080, lr 0.001, loss 0.214, epe 1.38, d1 10.4% epoch 160, lr 0.001, loss 0.177, epe 0.98, d1 6.1% epoch 240, lr 0.0001, loss 0.164, epe 0.89, d1 4.8% epoch 320, lr 0.0001, loss 0.158, epe 0.84, d1 4.2%

如果你看到loss一直在降但验证集指标不动，可能是过拟合了，需要加一些正则化或者数据增强。如果loss直接变成NaN，大概率是数值稳定性的问题，可以把输入图像归一化到0到1之间，或者降低学习率。

模型保存务必使用checkpoint方式，把model、optimizer、scheduler、epoch等信息都存下来，方便断点续训：

torch.save({ 'epoch': epoch, 'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), 'scheduler_state_dict': scheduler.state_dict(), }, 'checkpoint_latest.pth')

加载checkpoint时，要注意模型是否用了DataParallel或DistributedDataParallel，如果保存时带module.前缀，加载到普通模型上会报key不匹配的错，需要把前缀去掉：

state_dict = torch.load('checkpoint_latest.pth')['model_state_dict'] from collections import OrderedDict new_state_dict = OrderedDict() for k, v in state_dict.items(): if k.startswith('module.'): k = k[7:] new_state_dict[k] = v model.load_state_dict(new_state_dict)

5. 常见问题排查与避坑技巧实录

5.1 报错速查表

整个复现过程中，我遇到过的典型报错整理成了一张表，方便快速定位问题：

5.2 我反复踩的几个隐性坑

除了上面表格里的问题，还有几个隐性但影响很大的坑。第一个是验证时忘记把BatchNorm层切换到eval模式。BatchNorm在训练和推理时的行为不一样，如果忘了model.eval()，验证集指标会高得离谱，而且每次验证结果都不一样。我当时就因为这个，以为模型训练崩了，排查了半天才发现是这个小问题。

第二个坑是验证集划分。KITTI training只有200对图像，如果你在__getitem__里直接用整个目录做训练，没有单独留验证集，最终打印的“验证”指标其实是在训练集上算的，看起来很好但实际泛化能力未知。正确做法是先把文件列表切分出训练和验证两部分，再分别构造Dataset。

第三个是评估时的mask选择。如果用了disp_occ_0做训练，但评估时也直接用disp_occ_0并只过滤0值像素，那么D1指标会偏高，因为这些像素本来不需要被评估。KITTI官方评估通常分为非遮挡区域和全部区域两个指标。做复现对比时，统一一个mask逻辑，千万别今天用noc明天用occ，最后得到的结果完全没有可比性。

5.3 一些加快实验节奏的调试技巧

最后分享几个我后来逐步养成的实验习惯。

先跑通小规模再上全量。我第一次复现时，直接把全部200对图像、400个epoch一次跑满，结果第三天发现代码里有个bug，浪费了大量时间。后来我学乖了：先用10对图像、10个epoch跑一遍，只要能正常出loss、指标在下滑，再切回全量训练。这个冒烟测试看起来多花半小时，实际上能省好几天。

固定随机种子。PyTorch的DataLoader、模型初始化、数据增强都有随机性，如果不固定种子，可能连续两次训练得到的结果差异很大，干扰你判断代码改动是否有效。建议在训练脚本开头加上：

import random import numpy as np import torch def set_seed(seed): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed)

及时可视化视差图。训练到一半，把预测的视差图保存下来，用伪色彩图看一眼。如果预测图边缘有大量黑色条纹，可能是mask或上采样处理不对。如果整张图颜色很淡，可能是soft argmin后视差范围偏小。可视化能帮你快速发现很多loss指标看不出来的问题。

把实验配置写进文件名或json配置里。模型结构、超参数、数据增强策略、checkpoint路径，这些信息最好和日志、模型权重一起保存下来。我自己就吃过亏，两周后回看一个模型的权重，完全想不起来当时用了什么配置，只能重新跑实验对比。

我在实际使用中发现，复现PSMNet最大的收获倒不是最终指标有多好看，而是整个pipeline里每个环节都亲手摸了一遍。环境配置让你学会管理深度学习环境，数据读取让你理解KITTI的存储格式，训练调试让你真正明白损失函数和评估指标的关系。等这套流程跑通，后面想换新的立体匹配模型，基本上只需要改网络结构部分，其他环节都能复用。最后再分享一个小技巧：把KITTI的读取、评估、可视化写成独立模块，一开始就同时跑起来，调试网络结构时会轻松非常多。