⚡

FoundationStereo立体深度估计Skill tao-train-foundation-stereo

该技能围绕NVIDIA TAO FoundationStereo模型，提供立体匹配深度估计（视差估计）的完整工作流，包括数据集格式转换、模型训练、评估、推理、导出和量化，支持Middlebury、KITTI、ETH3D等多种数据集，适用于3D重建、自动驾驶感知等任务。关键词：立体匹配、视差估计、深度估计、FoundationStereo、3D重建、TAO深度网络、模型训练、评估、推理、导出。

视觉模型训练 0 次安装 0 次浏览 更新于 9/6/2026
名称 tao-train-foundation-stereo
描述 使用FoundationStereo进行立体深度估计。从立体图像对预测视差图，用于3D重建。在训练、评估、导出或对TAO FoundationStereo模型运行推理时使用。触发短语包括"train stereo depth"、“FoundationStereo”、“stereo disparity estimation”、“3D reconstruction from stereo”。
开源协议 Apache-2.0 compatibility: 需要docker + nvidia-container-toolkit。 metadata:
版本 “0.1.0”
作者 NVIDIA Corporation allowed-tools: Read Bash tags: - 立体 - 深度 - 估计

深度网络立体

独立安装？ 如果此会话未由TAO技能库插件初始化，请先运行tao-setup技能（主机预检、凭据、跨技能发现）。

使用FoundationStereo架构进行立体深度估计。从立体图像对预测视差图，用于3D重建。

使用预训练的Depth Anything v2和EdgeNeXt编码器。设置model.stereo_backbone.depth_anything_v2_pretrained_path和model.stereo_backbone.edgenext_pretrained_path。

单目和立体技能均调用容器内的统一TAO depth_net CLI；通过model.model_type（例如FoundationStereo）选择单目/立体家族。

此模型技能打包的PyT操作：train、evaluate、inference、export和quantize。PyT depth_net入口点在当前TAO镜像中不接受gen_trt_engine操作；仅通过部署工作流构建TensorRT引擎。

对于TAO部署TensorRT操作（gen_trt_engine、TensorRT evaluate和TensorRT inference），首先阅读references/tao-deploy-foundation-stereo.md。部署规范模板位于此技能的references/spec_template_deploy.yaml。

训练操作策略

此模型在模型层启用AutoML。在处理任何训练阶段请求之前，先阅读references/skill_info.yaml，并从显式automl_policy值或用户的工作流请求中解决运行覆盖。默认使用automl_policy: on，并且只在新启动提示中暴露on/off。将“关闭AutoML”、“禁用AutoML”、“无HPO”或“普通训练”等短语视为本次运行的automl_policy: off。当automl_policy: on、automl_enabled: true，并且同时打包了schemas/train.schema.json和references/spec_template_train.yaml时，默认通过tao-skill-bank:tao-run-automl路由训练操作，并带上此模型的skill_dir。保留工作流/应用程序对数据集、规范、输出目录、GPU/平台设置、父检查点和automl_policy的覆盖。仅当automl_policy: off或打包的训练模式/模板缺失时，才直接使用模型训练；在缺失模式的情况下，报告AutoML已启用但此模型不可运行，直到生成模式。

非训练操作如evaluate、inference、export和部署流程留在此模型技能中。每次运行的automl_policy覆盖不会更改模型元数据。

工作流

先决条件——数据可达性

您的数据集（左+右图像+GT视差）必须从容器内可访问：

  • SDK运行器：将文件放在运行器解析的S3路径（规范覆盖中显示的S3_TRAIN/S3_EVAL占位符）。运行器处理S3→容器路径挂载。
  • 直接docker run（例如本地测试）：将主机数据集根目录只读挂载到相同的容器内路径：
docker run ... -v <host_data_root>:<host_data_root>:ro <container> ...

相同的可达性要求适用于所有操作写入的<output_dir>。

步骤1 — 注释文件

由data_sources[*].data_file引用的逐行注释文件：

列数 格式 用途
2 <left> <right> 立体推理（无GT）
3 <left> <right> <disparity> 立体重建（带GT）
4 <left> <right> <disparity> <occlusion_mask> 立体重建（带GT和遮挡掩码）

如果您已有，直接指向它。否则通过depth_net convert生成：

depth_net convert -e <convert_spec.yaml>

convert_spec.yaml模板（立体）：

results_dir: <写入生成的注释文件的目录>
data_root: <其直接子级为包含图像+深度文件的场景文件夹的目录；convert递归遍历data_root，但期望子目录位于下方一层>
image_dir_pattern: [<匹配左图像路径的子串>]
right_dir_pattern: [<匹配右图像路径的子串>]
depth_dir_pattern: [<匹配GT视差路径的子串>]
nocc_dir_pattern: []                 # 可选，遮挡掩码路径
image_extension: '.png'  # 始终包含前导点
depth_extension: '.png'  # 格式必须与image_extension一致（替换只是子字符串替换）
nocc_extension: ''
split_ratio: 0.0        # 0.0/1.0 = 仅测试；0.8 = 80/20训练+验证

convert递归遍历data_root，选择路径字符串包含image_dir_pattern中所有子串的路径（AND过滤器），然后通过将image_dir_pattern[0]替换为对应模式的第一个元素并交换扩展名来派生右/深度/掩码路径。检查数据集的目录结构，并确定区分左、右和GT的子字符串（例如Middlebury的im0与im1与disp0GT）。

步骤2 — 根据数据配对model_type和dataset_name

当您的布局与支持的数据集类匹配时，优先使用特定于数据集的类——它会应用特定于类的路径约定、评估裁剪，以及（适用时）遮挡掩码处理。仅对于与任何注册类不匹配的布局，回退到GenericDataset。

数据类别 model_type dataset_name
Middlebury数据 FoundationStereo Middlebury
KITTI数据 FoundationStereo Kitti
ETH3D数据 FoundationStereo Eth3d
FSD合成数据 FoundationStereo FSD
IsaacReal合成数据 FoundationStereo IsaacRealDataset
Crestereo合成数据 FoundationStereo Crestereo
其他/非规范布局 FoundationStereo GenericDataset

立体data_sources的有效dataset_name值（不区分大小写）：FSD、IsaacRealDataset、Crestereo、Middlebury、Eth3d、Kitti、GenericDataset。

相同的dataset_name值适用于训练和评估操作（所有这些操作都使用带有GT视差的3列或4列注释）。部署侧的evaluate操作遵循相同规则——参见references/tao-deploy-foundation-stereo.md。对于使用2列注释（左+右，无GT）的推理，无论数据布局如何，都使用dataset_name: GenericDataset——数据集特定类（Middlebury/Kitti/Eth3d/FSD/IsaacRealDataset/Crestereo）需要3列输入，并在数据加载器级别拒绝2列注释。对于使用3列注释（左+右+GT）的推理，数据集特定类没问题。

步骤3 — 从规范覆盖编写spec yaml

从references/spec-overrides-foundation-stereo.md复制操作块。替换：

  • 步骤2中的model.model_type（通常为FoundationStereo）
  • 步骤2中的dataset.<...>.data_sources[*].dataset_name
  • 步骤1中的dataset.<...>.data_sources[*].data_file
  • 对于部署侧evaluate：强制dataset.test_dataset.batch_size: 1（参见references/tao-deploy-foundation-stereo.md）。

形状一致性：dataset.test_dataset.augmentation.crop_size中的crop_size应匹配export.input_height/input_width，以便训练模型评估器和部署侧TensorRT评估器在相同形状下运行。注意，crop_size在pyt evaluate路径上是装饰性的，但在部署evaluate侧是权威的——参见references/troubleshooting-foundation-stereo.md和references/tao-deploy-foundation-stereo.md。

全新安装的烟雾测试在crop_size: [128, 128]、dataset.max_disparity: 128和model.max_disparity: 128下验证。避免112×112裁剪，避免为烟雾测试将max_disparity设置为小于方形裁剪边长：这些组合可能在生成检查点之前，在FoundationStereo内部因特征图或损失掩码形状不匹配而失败。

数据源覆盖对于每个操作都是强制性的。每个data_sources条目是一个字典，具有两个必填字段：data_file和dataset_name。参见references/spec-overrides-foundation-stereo.md获取每个操作的数据集需求表、每个操作的覆盖块以及quantize已知问题说明。

步骤4 — 运行

在使用--user之前，在挂载的输出路径内创建可写的home/cache目录。一些TAO容器没有主机UID的/etc/passwd条目，当以该UID运行时，PyTorch/matplotlib需要可写的缓存路径。

mkdir -p <output_dir>/home \
         <output_dir>/.cache/matplotlib \
         <output_dir>/.cache/torchinductor \
         <output_dir>/.cache/xdg

docker run --gpus 'device=0' --shm-size 16G --ipc=host \
  --user "$(id -u):$(id -g)" \
  -e USER="$(id -un)" \
  -e LOGNAME="$(id -un)" \
  -e HOME=<output_dir>/home \
  -e MPLCONFIGDIR=<output_dir>/.cache/matplotlib \
  -e TORCHINDUCTOR_CACHE_DIR=<output_dir>/.cache/torchinductor \
  -e XDG_CACHE_HOME=<output_dir>/.cache/xdg \
  -v <data_root>:<data_root>:ro \
  -v <output_dir>:<output_dir> \
  <container> \
  depth_net <action> -e <spec.yaml>

如果没有--user "$(id -u):$(id -g)"，容器会以nobody:nogroup身份写入输出，阻止主机侧清理/重试。

步骤5 — 验证

  • 容器退出代码0
  • status.json中的kpi块已填充
  • 对于train：直接检查逐步骤train_loss（即使损失为NaN，入口点也报告Execution status: PASS）
  • 对于evaluate：依赖epe/bp1/bp2/bp3/d1/rmse（评估器也输出abs_rel/sq_rel/rmse_log，这些对立体不具意义——参见references/parameters-foundation-stereo.md）
  • 对于inference：results_dir下的工件

对于TAO部署TensorRT操作（gen_trt_engine、TensorRT evaluate和TensorRT inference），首先阅读references/tao-deploy-foundation-stereo.md。部署规范模板位于此技能的references/文件夹中，前缀为spec_template_deploy_*.yaml。

训练要求

  • 监控指标： val/loss
  • 评估数据集： 可选。通过dataset.val_dataset.data_sources配置验证数据集（每个条目需要data_file和dataset_name）。

参见references/spec-overrides-foundation-stereo.md获取每个操作的数据集需求表和每个操作必填的数据源覆盖块。

参数、指标、多GPU、导出/TRT、硬件

参见references/parameters-foundation-stereo.md获取完整的重要参数列表（包括model.encoder vits覆盖、model.max_disparity默认416、model.volume_dim无效说明、dataset.baseline、dataset.focal_x、train.precision、export.batch_size）、评估指标表、多GPU/多节点启动键、导出/TRT默认值（opset_version/on_cpu配对、NGC 576×960设置）和硬件要求。

错误模式与故障排除

参见references/troubleshooting-foundation-stereo.md了解视差溢出、烟雾测试形状不匹配、缺失预训练路径、encoder/dataset_name结构错误、depth_net_stereo: not found入口点说明、pyt与部署crop_size讨论，以及部署evaluate标量转换失败。

Spec参数/父模型推理

参见references/checkpoint-inference-mappings-foundation-stereo.md了解检查点解析规则（model_epoch_<epoch>_step_<step>.pth、dn_model_latest.pth策略）、父PyT gen_trt_engine的缺失，以及从depth_net_stereo.config.json生成的完整每操作推理映射表（包括parent_model/parent_job_id解析）。

部署