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.

            

            
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
        

tensorrt

The tensorrt parameter defines TensorRT engine generation.

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

            

            
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:

            

            
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:

            

            
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:

            

            
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:

            

            
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:

            

            
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:

            

            
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:

            

            
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.