Mask2former with TAO Deploy#

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

Each task is explained in detail in the following sections.

Converting .onnx File into TensorRT Engine#

You can reuse the spec from the Mask2Former export configuration 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: 960
  input_height: 544
  tensorrt:
    data_type: fp16
    workspace_size: 1024
    min_batch_size: 1
    opt_batch_size: 10
    max_batch_size: 10
    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	960	The input width	>0
`input_height`	unsigned int	544	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
`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 Mask2former from the exported ONNX at
``s3://my-bucket/mask2former/model.onnx`` using ``trt-spec.yaml``. Write
the engine to ``s3://my-bucket/mask2former/model.engine``. Run on the local Docker backend.

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
data:
  type: 'coco_panoptic'
  val:
    name: "coco_2017_val_panoptic"
    panoptic_json: "/datasets/coco/annotations/panoptic_val2017.json"
    img_dir: "/datasets/coco/val2017"
    panoptic_dir: "/datasets/coco/panoptic_val2017"
    batch_size: 1
    num_workers: 2

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

Evaluate the Mask2former TensorRT engine at
``s3://my-bucket/mask2former/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
  color_map: /path/to/colors.yaml
  label_map: /path/to/labels.csv
data:
  type: 'coco_panoptic'
  test:
    img_dir: /path/to/test_images/
    batch_size: 1

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

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

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