MAE with TAO Deploy

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

Each task is explained in detail in the following sections.

Converting MAE .onnx File into TensorRT Engine

You can reuse the spec from Exporting the model 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
  input_channel: 3
  input_width: 224
  input_height: 224
  tensorrt:
    data_type: fp16
    workspace_size: 1024
    min_batch_size: 1
    opt_batch_size: 16
    max_batch_size: 16

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	224	The input width.	>0
input_height	unsigned int	224	The input height.	>0
batch_size	unsigned int	-1	The batch size of the ONNX model.	>=-1
verbose	bool	False	Enables verbosity for the TensorRT log.

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

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

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