Stereo Depth with TAO Deploy#

To generate an optimized NVIDIA® TensorRT™ engine for a stereo depth model, the gen_trt_engine action takes an ONNX file previously produced by the corresponding export action. TAO Deploy supports two stereo depth networks:

  • FoundationStereo — the original high-accuracy stereo depth network.

  • FastFoundationStereo (FFS) — a distilled, lower-latency variant introduced in TAO 7.0.1.

For more information about training these models, refer to Stereo Depth Estimation.

Note

The deploy network is selected by the model.model_type field in the spec file: FoundationStereo or FastFoundationStereo. Both networks share the same gen_trt_engine, evaluate, and inference deploy actions and command syntax; only the spec contents differ. The main differences for FastFoundationStereo are summarized in FastFoundationStereo notes.

FastFoundationStereo Notes#

FastFoundationStereo reuses the FoundationStereo deploy path verbatim — the same export → gen_trt_engine → evaluate → inference workflow and the same spec fields apply. Keep the following FFS-specific differences in mind:

  • Disparity range. model.max_disparity defaults to 416 (the FoundationStereo range). The FastFoundationStereo commercial checkpoint is trained with a different range (for example 192 for the bp2 checkpoint), so this field must be set to match the checkpoint. The deploy evaluator treats model.max_disparity as authoritative because FastFoundationStereo and FoundationStereo can use different ranges; an incorrect value shifts disparity out of the trained regime and degrades metrics. Set the same value used at training and export time.

  • GWC feature normalization. model.gwc_feature_normalize (default True) is consumed only by FastFoundationStereo. The FoundationStereo path ignores it.

  • Architecture overrides. The FastFoundationStereo checkpoint requires its full set of width/architecture overrides (for example hidden_dims, n_gru_layers, valid_iters, volume_dim, and the distilled width fields) to match the trained weights. Reuse the FFS experiment/export spec as the starting point so these fields are consistent across export and gen_trt_engine. The schema defaults are hidden_dims: [128, 128, 128], n_gru_layers: 3, valid_iters: 22, and volume_dim: 32; these are shown in the FastFoundationStereo gen_trt_engine example below.

Converting ONNX File into TensorRT Engine#

You can reuse the spec from the FoundationStereo export action as a starting point.

gen_trt_engine#

The gen_trt_engine parameter defines TensorRT engine generation.

gen_trt_engine:
  onnx_file: /path/to/onnx_file
  trt_engine: /path/to/trt_engine
  batch_size: -1
  tensorrt:
    data_type: fp16
    workspace_size: 1024
    min_batch_size: 1
    opt_batch_size: 2
    max_batch_size: 4

Field

value_type

Description

default_value

valid_min

valid_max

valid_options

automl_enabled

results_dir

string

Path to where all the assets generated from a task are stored.

FALSE

gpu_id

int

Index of the GPU to build the TensorRT engine.

0

FALSE

onnx_file

string

Path to the ONNX model file.

???

FALSE

trt_engine

string

Path where the generated TensorRT engine from gen_trt_engine is stored. This only works with tao-deploy.

FALSE

input_channel

int

Number of channels in the input tensor.

3

3

FALSE

opset_version

int

Operator set version of the ONNX model used to generate the TensorRT engine.

17

1

FALSE

batch_size

int

Batch size of the input tensor for the engine. A value of -1 implies dynamic tensor shapes.

-1

-1

FALSE

verbose

bool

Flag to enable verbose TensorRT logging.

False

FALSE

tensorrt

collection

Hyperparameters to configure the TensorRT Engine builder.

FALSE

tensorrt#

The tensorrt parameter defines the TensorRT engine generation.

Field

value_type

Description

default_value

valid_min

valid_max

valid_options

automl_enabled

data_type

string

Precision to be set for building the TensorRT engine.

FP32

FP32,FP16, BF16

FALSE

workspace_size

int

Size in megabytes of the workspace TensorRT has to run its optimization tactics and generate the TensorRT engine.

1024

FALSE

min_batch_size

int

Minimum batch size in the optimization profile for the input tensor of the TensorRT engine.

1

FALSE

opt_batch_size

int

Optimum batch size in the optimization profile for the input tensor of the TensorRT engine.

1

FALSE

max_batch_size

int

Maximum batch size in the optimization profile for the input tensor of the TensorRT engine.

1

FALSE

For FastFoundationStereo, use the same gen_trt_engine action with a spec that selects FastFoundationStereo and carries the matching model fields. A static-shape FP16 engine is the production-validated path:

model:
  model_type: FastFoundationStereo
  max_disparity: 192        # must match the FastFoundationStereo checkpoint
  gwc_feature_normalize: true   # default; consumed only by FastFoundationStereo
  # Architecture overrides must match the trained FFS weights. The values below
  # are the schema defaults; replace them with the values from your FFS
  # experiment/export spec when they differ.
  hidden_dims: [128, 128, 128]
  n_gru_layers: 3
  valid_iters: 22
  volume_dim: 32
gen_trt_engine:
  onnx_file: /path/to/onnx_file
  trt_engine: /path/to/trt_engine
  batch_size: 1             # static shapes; set -1 for a dynamic-batch profile
  tensorrt:
    data_type: fp16
    workspace_size: 4096    # FFS typically needs more than the 1024 default
    min_batch_size: 1
    opt_batch_size: 1
    max_batch_size: 1

Ask the agent to run the gen_trt_engine action against your spec. For example:

Build an FP16 TensorRT engine for FoundationStereo from the exported
ONNX at ``s3://my-bucket/stereo/model.onnx`` using ``trt-spec.yaml``.
Write the engine to ``s3://my-bucket/stereo/model.engine``. Run on
the local Docker daemon.

The same request works for FastFoundationStereo by pointing at the FFS ONNX and spec:

Build a static-shape FP16 TensorRT engine for FastFoundationStereo from the
exported ONNX at ``s3://my-bucket/stereo/ffs.onnx`` using ``ffs-trt-spec.yaml``.
Write the engine to ``s3://my-bucket/stereo/ffs.engine``. Run on the local
Docker daemon.

Running Evaluation through a TensorRT Engine#

You can reuse the TAO evaluation specification file for evaluation through a TensorRT engine. The following is a sample specification file:

evaluate:
  trt_engine: /path/to/engine/file
  input_width: 736
  input_height: 320
dataset:
  dataset_name: StereoDataset
  test_dataset:
    data_sources:
      - dataset_name: GenericDataset
        data_file: /data/depth_net/annotations_test.txt
    batch_size: 4
    workers: 4

For FastFoundationStereo, reuse the same evaluate action and spec format. Set model.model_type: FastFoundationStereo and model.max_disparity to the value used during training/export so the evaluator masks and scores disparities against the correct range:

model:
  model_type: FastFoundationStereo
  max_disparity: 192
evaluate:
  trt_engine: /path/to/engine/file
  input_width: 736
  input_height: 480
dataset:
  dataset_name: StereoDataset
  max_disparity: 192
  test_dataset:
    data_sources:
      - dataset_name: GenericDataset
        data_file: /data/depth_net/annotations_test.txt
    batch_size: 1
    workers: 4

Ask the agent to run the evaluate action against the engine you built. For example:

Evaluate the FoundationStereo TensorRT engine at
``s3://my-bucket/stereo/model.engine`` against ``eval-spec.yaml``. Run on
local Docker.

The same applies to FastFoundationStereo:

Evaluate the FastFoundationStereo TensorRT engine at
``s3://my-bucket/stereo/ffs.engine`` against ``ffs-eval-spec.yaml``. Run on
local Docker.

Running Inference through a TensorRT Engine#

You can reuse the TAO inference specification file for inference through a TensorRT engine. This is a sample specification file:

inference:
  input_width: 736
  input_height: 320
  trt_engine: /path/to/engine/file
dataset:
  dataset_name: StereoDataset
  infer_dataset:
    data_sources:
      - dataset_name: GenericDataset
        data_file: /data/depth_net/annotations_test.txt
  workers: 4
  batch_size: 4

For FastFoundationStereo, reuse the same inference action and spec format with model.model_type: FastFoundationStereo and the matching model.max_disparity:

model:
  model_type: FastFoundationStereo
  max_disparity: 192
inference:
  input_width: 736
  input_height: 480
  trt_engine: /path/to/engine/file
dataset:
  dataset_name: StereoDataset
  infer_dataset:
    data_sources:
      - dataset_name: GenericDataset
        data_file: /data/depth_net/annotations_test.txt
  workers: 4
  batch_size: 1

Ask the agent to run the inference action against the engine you built. For example:

Run FoundationStereo inference with the TensorRT engine at
``s3://my-bucket/stereo/model.engine`` using ``infer-spec.yaml``. Run on
your chosen backend.

The same request works for FastFoundationStereo:

Run FastFoundationStereo inference with the TensorRT engine at
``s3://my-bucket/stereo/ffs.engine`` using ``ffs-infer-spec.yaml``. Run on
your chosen backend.

Annotated visualizations are written to images_annotated under the configured results directory, and predictions are written to labels.

Note

TensorRT inference writes colorized disparity visualizations (.png). Raw .pfm disparity output (the inference.save_raw_pfm field) is only produced by the PyTorch inference path, not by the TAO Deploy engine path.