Stereo Depth with TAO Deploy#

To generate an optimized NVIDIA® TensorRT™ engine for FoundationStereo, the gen_trt_engine action takes an ONNX file previously produced by the FoundationStereo export action. For more information about training a FoundationStereo model, refer to Stereo Depth Estimation.

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

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

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.

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

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.

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

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.

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