Segformer with TAO Deploy

Segformer etlt file generated from tao export is taken as an input to tao-deploy to generate optimized TensorRT engine. For 4.0.0, we do not support Int8 precision for Segformer.

Converting .etlt File into TensorRT Engine

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

trt_config

The trt_config parameter provides options related to TensorRT generation.

Copy
Copied!

            
            trt_config:
  data_type: FP32
  workspace_size: 1024
  min_batch_size: 1
  opt_batch_size: 1
  max_batch_size: 1

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 optimization profile shape	>0
`opt_batch_size`	unsigned int	1	The optimal batch size used for optimization profile shape	>0
`max_batch_size`	unsigned int	1	The maximum batch size used for 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 \
           encryption_key=<key> \
           output_file=/path/to/etlt/file \
           trt_engine=/path/to/engine/file \
           trt_config.data_type=<data_type>

Required Arguments

-e, --experiment_spec: The experiment spec file to set up the TensorRT engine generation. This should be the same as the export specification file.
-k, --key: A user-specific encoding key to load a .etlt model.
output_file: The .etlt model to be converted.
trt_engine: The path where the generated engine will be stored.
data_type: Segformer only supports FP32 and FP16.

Sample Usage

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

Copy
Copied!

            
            tao-deploy segformer gen_trt_engine -e $DEFAULT_SPEC
           encryption_key=$YOUR_KEY \
           model_path=$ETLT_FILE \
           trt_engine=$ENGINE_FILE \
           trt_config.data_type=FP16

Running Evaluation through TensorRT Engine

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

Copy
Copied!

            
            exp_config:
    gpu_ids:
        - 0
    num_gpus: 4
    manual_seed: 49
    distributed: True
model_config:
    pretrained: null
    backbone:
        type: "mit_b1"
dataset_config:
    img_norm_cfg:
        mean:
            - 127.5
            - 127.5
            - 127.5
        std:
            - 127.5
            - 127.5
            - 127.5
test_pipeline:
    multi_scale:
    - 2048
    - 512
input_type: "grayscale"
data_root: /tlt-pytorch
test_img_dir: /data/images/val
test_ann_dir: /data/masks/val
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_per_gpu: 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 \
           model_path=/path/to/engine/file \
           output_dir=/path/to/outputs

Required Arguments

-e, --experiment_spec: The experiment spec file for evaluation. This should be the same as the tao evaluate specification file.
model_path: The engine file to run evaluation.
output_dir: The directory where evaluation results will be stored.

Sample Usage

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

Copy
Copied!

            
            tao-deploy segformer evaluate -e $DEFAULT_SPEC
           model_path=$ENGINE_FILE \
           output_dir=$RESULTS_DIR

Running Inference through TensorRT Engine

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

Copy
Copied!

            
            exp_config:
    gpu_ids:
        - 0
    num_gpus: 4
    manual_seed: 49
    distributed: True
model_config:
    pretrained: null
    backbone:
        type: "mit_b1"
dataset_config:
    img_norm_cfg:
        mean:
            - 127.5
            - 127.5
            - 127.5
        std:
            - 127.5
            - 127.5
            - 127.5
test_pipeline:
    multi_scale:
    - 2048
    - 512
input_type: "grayscale"
data_root: /tlt-pytorch
test_img_dir: /data/images/val
test_ann_dir: /data/masks/val
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
samples_per_gpu: 1
workers_per_gpu: 1

Use the following command to run Segformer engine inference:

Copy
Copied!

            
            tao-deploy segformer inference -e /path/to/spec.yaml \
           model_path=/path/to/engine/file \
           output_dir=/path/to/outputs

Required Arguments

-e, --experiment_spec: The experiment spec file for inference. This should be the same as the tao inference specification file.
model_path: The engine file to run inference.
output_dir: The directory where inference results will be stored.

Sample Usage

Here’s an example of using the inference command to run inference with the TensorRT engine:

Copy
Copied!

            
            tao-deploy segformer inference -e $DEFAULT_SPEC
           model_path=$ENGINE_FILE \
           output_dir=$RESULTS_DIR

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.