CenterPose with TAO Deploy#

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

Converting an ONNX File into TensorRT Engine#

The gen_trt_engine section of the spec configures TensorRT engine generation. You can reuse the spec from the CenterPose export action as a starting point.

gen_trt_engine:
onnx_file: /path/to/onnx_file
trt_engine: /path/to/trt_engine
input_channel: 3
input_width: 512
input_height: 512
tensorrt:
  data_type: fp32
  workspace_size: 1024
  min_batch_size: 1
  opt_batch_size: 2
  max_batch_size: 4
  calibration:
    cal_image_dir: /path/to/cal/images
    cal_cache_file: /path/to/cal.bin
    cal_batch_size: 10
    cal_batches: 1000

Parameter	Datatype	Default	Description	Supported Values
`onnx_file`	string		The precision to be used for the TensorRT engine
`trt_engine`	string		The maximum workspace size for the TensorRT engine
`input_channel`	unsigned int	3	The input channel size. Only the value 3 is supported.	3
`input_width`	unsigned int	512	The input width	>0
`input_height`	unsigned int	512	The input height	>0
`batch_size`	unsigned int	-1	The batch size of the ONNX model	>=-1

tensorrt#

The tensorrt parameter defines TensorRT engine generation.

Parameter	Datatype	Default	Description	Supported Values
`data_type`	string	fp32	The precision to be used for the TensorRT engine	fp32/fp16/int8
`workspace_size`	unsigned int	1024	The maximum workspace size for the TensorRT engine	>1024
`min_batch_size`	unsigned int	1	The minimum batch size used for the optimization profile shape	>0
`opt_batch_size`	unsigned int	1	The optimal batch size used for the optimization profile shape	>0
`max_batch_size`	unsigned int	1	The maximum batch size used for the optimization profile shape	>0

calibration#

The calibration parameter defines TensorRT engine generation with PTQ INT8 calibration.

Parameter	Datatype	Default	Description	Supported Values
`cal_image_dir`	string		The list of paths that contain images used for calibration
`cal_cache_file`	string		The path to the calibration cache file to be dumped
`cal_batch_size`	unsigned int	1	The batch size per batch during calibration	>0
`cal_batches`	unsigned int	1	The number of batches to calibrate	>0

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

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

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
  opencv: False
  eval_num_symmetry: 1
  results_dir: /path/to/save/results
dataset:
  test_data: /path/to/testing/images/and/json/files
  batch_size: 2
  workers: 4

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

Evaluate the CenterPose TensorRT engine at
``s3://my-bucket/centerpose/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. The following is a sample specification file:

inference:
  trt_engine: /path/to/engine/file
  visualization_threshold: 0.3
  principle_point_x: 298.3
  principle_point_y: 392.1
  focal_length_x: 651.2
  focal_length_y: 651.2
  skew: 0.0
  axis_size: 0.5
  use_pnp: True
  save_json: True
  save_visualization: True
  opencv: True
dataset:
  inference_data: /path/to/inference/files
  batch_size: 1
  workers: 4

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

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

Visualization results are written to the configured results directory.