Classification (PyTorch) with TAO Deploy

To generate an optimized TensorRT engine, a classification (PyTorch) .onnx file, which is first generated using tao model classification_pyt export, is taken as an input to tao deploy classification_pyt gen_trt_engine. For more information about training a classification (PyTorch) model, refer to the Classification PyTorch training documentation. With TAO 5.0.0, INT8 precision is not supported for classification (PyTorch) models.

To convert the .onnx file, you can reuse the spec file from the tao model classification_pyt export command.

gen_trt_engine

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_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

Use the following command to run classification (PyTorch) engine generation:

Copy
Copied!
            

tao deploy classification_pyt gen_trt_engine -e /path/to/spec.yaml \ -r /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

  • -r, --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 classification_pyt 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


You can reuse the TAO evaluation spec file for evaluation through a TensorRT engine. The classes field is only required if you are using a custom class names. If this field is not provided, class mapping is based on the alphanumerical order of the image folder names. The following is a sample spec file:

Copy
Copied!
            

evaluate: trt_engine: /path/to/engine/file topk: 1 dataset: data: samples_per_gpu: 16 test: data_prefix: /raid/ImageNet2012/ImageNet2012/val classes: /raid/ImageNet2012/classnames.txt

Use the following command to run classification (PyTorch) engine evaluation:

Copy
Copied!
            

tao deploy classification_pyt evaluate -e /path/to/spec.yaml \ -r /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

  • -r, --results_dir: The directory where the JSON status-log 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 classification_pyt evaluate -e $DEFAULT_SPEC -r $RESULTS_DIR \ evaluate.trt_engine=$ENGINE_FILE


You can reuse the TAO inference spec file for inference through a TensorRT engine. The following is a sample spec file:

Copy
Copied!
            

inference: trt_engine: /path/to/engine/file dataset: data: samples_per_gpu: 16 test: data_prefix: /raid/ImageNet2012/ImageNet2012/val classes: /raid/ImageNet2012/classnames.txt

Use the following command to run classification (PyTorch) engine inference:

Copy
Copied!
            

tao deploy classification_pyt inference -e /path/to/spec.yaml \ -r /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

  • -r, --results_dir: The directory where the JSON status-log file and inference results will be dumped

  • inference.trt_engine: The engine file for inference

Sample Usage

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

Copy
Copied!
            

tao deploy classification_pyt inference -e $DEFAULT_SPEC -r $RESULTS_DIR \ evaluate.trt_engine=$ENGINE_FILE

The visualization will be stored in $RESULTS_DIR/images_annotated, and the KITTI format predictions will be stored under $RESULTS_DIR/labels.

© Copyright 2023, NVIDIA.. Last updated on Dec 8, 2023.