UNet with TAO Deploy

To generate an optimized TensorRT engine for UNet, the gen_trt_engine action takes an ONNX file previously produced by the UNet export action. For more information about training the UNet, refer to the UNet training documentation.

Converting .onnx File into TensorRT Engine

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

Note

When generating a TensorRT engine for a model trained with QAT enabled, the tensor scale factors defined by the calibration cache file are required. However, the current version of QAT does not natively support DLA int8 deployment on Jetson. To deploy this model on a Jetson with DLA int8, force post-training quantization to generate the calibration cache file.

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

Build an INT8 TensorRT engine for UNet from the exported ONNX at
``s3://my-bucket/unet/unet.onnx`` using ``trt-spec.yaml``. Calibrate
against ``s3://my-bucket/unet/cal-images/`` and write the engine to
``s3://my-bucket/unet/int8.engine``. Run on the local Docker backend.

Running Evaluation through TensorRT Engine

Use the same specification file as the TAO evaluation specification file. The following is a sample specification file:

model_config {
  model_input_width: 320
  model_input_height: 320
  model_input_channels: 1
  num_layers: 18
  all_projections: true
  load_graph: true
  arch: "resnet"
  use_batch_norm: False
  training_precision {
    backend_floatx: FLOAT32
  }
}
dataset_config {
  dataset: "custom"
  augment: False
  input_image_type: "grayscale"
  train_images_path:"/workspace/tao-experiments/data/isbi/images/train"
  train_masks_path:"/workspace/tao-experiments/data/isbi/masks/train"
  val_images_path:"/workspace/tao-experiments/data/isbi/images/val"
  val_masks_path:"/workspace/tao-experiments/data/isbi/masks/val"
  test_images_path:"/workspace/tao-experiments/data/isbi/images/test"
  data_class_config {
    target_classes {
      name: "foreground"
      mapping_class: "foreground"
      label_id: 0
    }
    target_classes {
      name: "background"
      mapping_class: "background"
      label_id: 1
    }
  }
}

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

Evaluate the UNet TensorRT engine at ``s3://my-bucket/unet/int8.engine``
against ``eval-spec.yaml``. Run on local Docker.

Running Inference through TensorRT Engine

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

Run UNet inference with the TensorRT engine at
``s3://my-bucket/unet/int8.engine`` using ``infer-spec.yaml``. Run on
your chosen backend.

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