Segformer with TAO Deploy

Segformer ONNX file generated from tao export is taken as an input to tao-deploy to generate optimized TensorRT engine. We do not support Int8 precision for Segformer.

Converting .onnx File into TensorRT Engine

Same spec file can be used as the tao model segformer export command.

trt_config

The gen_trt_engine parameter defines TensorRT engine generation.

Copy
Copied!

            
            gen_trt_engine:
  onnx_file: /path/to/onnx_file
  trt_engine: /path/to/trt_engine
  input_width: 512
  input_height: 512
  tensorrt:
    data_type: FP32
    workspace_size: 1024
    min_batch_size: 1
    opt_batch_size: 1
    max_batch_size: 1

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

Use the following command to run Segformer engine generation:

Copy
Copied!

            
            tao deploy segformer gen_trt_engine -e /path/to/spec.yaml \
           results_dir=/path/to/results \
           gen_trt_engine.onnx_file=/path/to/onnx/file \
           gen_trt_engine.trt_engine=/path/to/engine/file \
           gen_trt_engine.tensorrt.data_type=<data_type>

Required Arguments

-e, --experiment_spec: The experiment spec file to set up TensorRT engine generation

Optional Arguments

results_dir: The directory where the JSON status-log file will be dumped
gen_trt_engine.onnx_file: The .onnx model to be converted
gen_trt_engine.trt_engine: The path where the generated engine will be stored
gen_trt_engine.tensorrt.data_type: The precision to be exported

Sample Usage

Here’s an example of using the gen_trt_engine command to generate an FP16 TensorRT engine:

Copy
Copied!

            
            tao deploy segformer gen_trt_engine -e $DEFAULT_SPEC
           gen_trt_engine.onnx_file=$ONNX_FILE \
           gen_trt_engine.trt_engine=$ENGINE_FILE \
           gen_trt_engine.tensorrt.data_type=FP16

Running Evaluation through TensorRT Engine

Same spec file as TAO evaluation/ inference spec file. Sample spec file:

Copy
Copied!

            
            model:
  input_height: 512
  input_width: 512
  backbone:
    type: "mit_b1"
dataset:
  img_norm_cfg:
    mean:
      - 127.5
      - 127.5
      - 127.5
    std:
      - 127.5
      - 127.5
      - 127.5
    test_dataset:
      img_dir: /data/images/val
      ann_dir: /data/masks/val
  input_type: "grayscale"
  data_root: /tlt-pytorch
  palette:
    - seg_class: foreground
      rgb:
        - 0
        - 0
        - 0
      label_id: 0
      mapping_class: foreground
    - seg_class: background
      rgb:
        - 255
        - 255
        - 255
      label_id: 1
      mapping_class: background
  batch_size: 1
  workers_per_gpu: 1

Use the following command to run Segformer engine evaluation:

Copy
Copied!

            
            tao deploy segformer evaluate -e /path/to/spec.yaml \
           results_dir=/path/to/results \
           evaluate.trt_engine=/path/to/engine/file

Required Arguments

-e, --experiment_spec: The experiment spec file for evaluation This should be the same as the tao evaluate spec file

Optional Arguments

results_dir: The directory where the evaluat file and evaluation results will be dumped
evaluate.trt_engine: The engine file for evaluation

Sample Usage

Here’s an example of using the evaluate command to run evaluation with a TensorRT engine:

Copy
Copied!

            
            tao deploy segformer evaluate -e $DEFAULT_SPEC
           results_dir=$RESULTS_DIR \
           evaluate.trt_engine=$ENGINE_FILE

Running Inference through TensorRT Engine

Use the following command to run SegFormer engine inference:

Copy
Copied!

            
            tao deploy segformer inference -e /path/to/spec.yaml \
           results_dir=/path/to/results \
           inference.trt_engine=/path/to/engine/file

Required Arguments

-e, --experiment_spec: The experiment spec file for inference. This should be the same as the tao inference spec file.

Optional Arguments

results_dir: The directory where JSON status-log file and inference results will be dumped
inference.trt_engine: The engine file for inference

Sample Usage

For inference, you can re-use the spec config mentioned under running_evaluation_through_tensorrt_engine. Here’s an example of using the inference command to run inference with a TensorRT engine:

Copy
Copied!

            
            tao deploy segformer inference -e $DEFAULT_SPEC
           results_dir=$RESULTS_DIR \
           inference.trt_engine=$ENGINE_FILE

The mask overlaid visualization will be stored under $RESULTS_DIR/vis_overlay and raw predictions in mask format will be stored under $RESULTS_DIR/mask_labels.