NVIDIA TAO Toolkit v4.0.1
NVIDIA TAO Release 4.0.1

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 Jul 27, 2023.