从KITTI到自动驾驶：用PyTorch复现PSMNet立体匹配网络（附完整训练代码）

立体视觉在自动驾驶领域扮演着关键角色，它通过分析左右摄像头图像的视差来估计场景深度。2018年CVPR会议上提出的PSMNet（Pyramid Stereo Matching Network）以其创新的金字塔池化模块和堆叠沙漏结构，在当时刷新了KITTI数据集的精度记录。本文将带你从零开始，在PyTorch框架下完整复现这一经典网络，并分享在实际自动驾驶项目中应用时的工程技巧。

1. 环境配置与数据准备

1.1 硬件选择与PyTorch环境搭建

根据我们的实测，在NVIDIA RTX 4090上训练PSMNet时，batch_size可设置为6而不会出现显存溢出。如果使用AutoDL等云平台，建议选择A100 40GB以上的实例。以下是快速搭建环境的命令：

bash复制conda create -n psmnet python=3.8
conda activate psmnet
pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 -f https://download.pytorch.org/whl/torch_stable.html
pip install opencv-python matplotlib tqdm tensorboard

注意：PyTorch 1.12版本在3D卷积运算上比新版本更稳定，这也是原论文采用的版本

1.2 KITTI数据集处理技巧

KITTI 2015数据集包含200组街景立体图像，但原始标注存在以下问题需要特别处理：

1. 稀疏标注补偿：激光雷达获取的视差图仅有约5%像素有标注，我们采用以下策略增强：

   python复制def sparse_to_dense(disparity, max_disp=192):
       """将稀疏视差图转换为密集视差图"""
       h, w = disparity.shape
       dense = np.zeros_like(disparity)
       # 使用最近邻插值填充无效区域
       valid_pixels = (disparity > 0) & (disparity < max_disp)
       yy, xx = np.where(valid_pixels)
       interp = NearestNDInterpolator(np.column_stack((yy, xx)), disparity[valid_pixels])
       dense = interp(np.indices(disparity.shape).T).T
       return dense
   

2. 动态裁剪策略：不同于原论文固定裁剪256×512，我们根据图像内容动态调整：

   python复制def dynamic_crop(left_img, right_img, disparity, crop_size=(256,512)):
       """根据有效视差区域智能裁剪"""
       valid_mask = (disparity > 0) & (disparity < max_disp)
       if valid_mask.sum() > 1000:
           y, x = np.where(valid_mask)
           center_y, center_x = int(y.mean()), int(x.mean())
       else:
           center_y, center_x = left_img.shape[0]//2, left_im