OCDNet with TAO Deploy
To generate an optimized TensorRT engine, the OCDNet .onnx
file, which is generated from tao export
, is
taken as an input to tao-deploy
. For more information about training an OCDNet model, refer to
the OCDNet training documentation.
gen_trt_engine
The gen_trt_engine
parameter in the experiment specification file provides options to generate a
TensorRT engine from the .onnx
file. The following is a sample config:
gen_trt_engine:
results_dir: /results/gen_trt_engine
onnx_file: /results/export/model_best.onnx
trt_engine: /results/export/ocdnet_model.engine
width: 1280
height: 736
img_mode: BGR
tensorrt:
data_type: int8
min_batch_size: 1
opt_batch_size: 1
max_batch_size: 1
calibration:
cal_image_dir: /data/ocdnet/train/img
cal_cache_file: /results/export/cal.bin
cal_batch_size: 8
cal_num_batches: 2
Parameter |
Datatype |
Default |
Description |
Supported Values |
|
string |
– |
The path to the results directory |
– |
|
string |
– |
The path to the |
– |
|
string |
– |
The absolute path to the generated TensorRT engine |
– |
|
unsigned int |
– |
The input width |
>0 |
|
unsigned int |
– |
The input height |
>0 |
|
string |
BGR |
The input image mode |
BGR,RGB,GRAY |
tensorrt
The tensorrt
parameter defines the TensorRT engine generation.
Parameter |
Datatype |
Default |
Description |
Supported Values |
|
string |
fp32 |
The precision to use for the TensorRT engine |
fp32/fp16/int8 |
|
unsigned int |
1024 |
The maximum workspace size for the TensorRT engine |
>1024 |
|
unsigned int |
1 |
The minimum batch size for the optimization profile shape |
>0 |
|
unsigned int |
1 |
The optimal batch size for the optimization profile shape |
>0 |
|
unsigned int |
1 |
The maximum batch size for the optimization profile shape |
>0 |
calibration
The calibration
parameter defines TensorRT engine generation with PTQ INT8 calibration.
Parameter |
Datatype |
Default |
Description |
Supported Values |
|
string list |
A list of paths that contain images used for calibration |
||
|
string |
The path to the calibration cache file to be dumped |
||
|
unsigned int |
1 |
The batch size per batch during calibration |
>0 |
|
unsigned int |
1 |
The number of batches to calibrate |
>0 |
Use the following command to run OCDNet engine generation:
tao deploy ocdnet gen_trt_engine -e /path/to/spec.yaml \
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_file
: The path to the experiment spec filegen_trt_engine.onnx_file
: The.onnx
model to be converted. This argument can be omitted if it is defined in spec file.gen_trt_engine.trt_engine
: The path where the generated engine will be stored. This argument can be omitted if it is defined in spec file.gen_trt_engine.tensorrt.data_type
: Precision to be exported. This argument can be omitted if it is defined in spec file.
Optional Arguments
-r, --results_dir
: The directory where the status log JSON file will be dumped
Sample Usage
Here’s an example of using the gen_trt_engine
command to generate an FP16 TensorRT engine:
tao deploy ocdnet gen_trt_engine -e $SPECS_DIR/gen_trt_engine.yaml \
gen_trt_engine.onnx_file=$RESULTS_DIR/export/model_best.onnx \
gen_trt_engine.trt_engine=$RESULTS_DIR/export/ocdnet_model.engine \
gen_trt_engine.tensorrt.data_type=fp16
Here’s an example of using the gen_trt_engine
command to generate an INT8 TensorRT engine:
tao deploy ocdnet gen_trt_engine -e $SPECS_DIR/gen_trt_engine.yaml \
gen_trt_engine.onnx_file=$RESULTS_DIR/export/model_best.onnx \
gen_trt_engine.tensorrt.min_batch_size=1 \
gen_trt_engine.tensorrt.opt_batch_size=3 \
gen_trt_engine.tensorrt.max_batch_size=3 \
gen_trt_engine.tensorrt.data_type=int8 \
gen_trt_engine.trt_engine=$RESULTS_DIR/export/ocdnet_model.engine
The TAO Evaluation spec file can be reused for this step. The following is a sample spec file:
model:
load_pruned_graph: False
pruned_graph_path: /results/prune/pruned_0.1.pth
evaluate:
results_dir: /results/evaluate
checkpoint: /results/train/model_best.pth
trt_engine: /results/export/ocdnet_model.engine
gpu_id: 0
post_processing:
type: SegDetectorRepresenter
args:
thresh: 0.3
box_thresh: 0.55
max_candidates: 1000
unclip_ratio: 1.5
metric:
type: QuadMetric
args:
is_output_polygon: false
dataset:
validate_dataset:
data_path: ['/data/ocdnet/test']
args:
pre_processes:
- type: Resize2D
args:
short_size:
- 1280
- 736
resize_text_polys: true
img_mode: BGR
filter_keys: []
ignore_tags: ['*', '###']
loader:
batch_size: 1
shuffle: false
pin_memory: false
num_workers: 4
Use the following command to run OCDNet engine evaluation:
tao deploy ocdnet 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 thetao evaluate
specification file.evaluate.trt_engine
: The engine file to run evaluation. It can be ignored if it is defined in the spec file.
Optional Arguments
-r, --results_dir
: The directory where status logging JSON file and evaluation results will be dumped.
Sample Usage
Here’s an example of using the evaluate
command to run evaluation with the TensorRT engine:
tao deploy ocdnet evaluate -e $SPECS_DIR/evaluate.yaml \
-r $RESULTS_DIR \
evaluate.trt_engine=$RESULTS_DIR/export/ocdnet_model.engine
You can reuse the TAO inference spec file for this step. The following is a sample spec file:
inference:
checkpoint: /results/train/model_best.pth
trt_engine: /results/export/ocdnet_model.engine
input_folder: /data/ocdnet/test/img
width: 1280
height: 736
img_mode: BGR
polygon: false
results_dir: /results/inference
post_processing:
type: SegDetectorRepresenter
args:
thresh: 0.3
box_thresh: 0.55
max_candidates: 1000
unclip_ratio: 1.5
Use the following command to run OCDNet engine inference:
tao deploy ocdnet inference -e /path/to/spec.yaml \
inference.trt_engine=$RESULTS_DIR/export/ocdnet_model.engine \
inference.input_folder=$DATA_DIR/test/img \
inference.results_dir=$RESULTS_DIR/inference
Required Arguments
-e, --experiment_spec_file
: The path to the experiment spec fileinference.trt_engine
: The engine file to run inference. This argument can be omitted if it is defined in spec file.inference.input_folder
: The engine file to run inference. This argument can be omitted if it is defined in spec file.inference.results_dir
: The directory where status logging JSON file and inference results will be dumped.
This argument can be omitted if it is defined in spec file.
Sample Usage
Here’s an example of using the inference
command to run inference with the TensorRT engine:
tao deploy ocdnet inference -e $SPECS_DIR/inference.yaml \
inference.trt_engine=$RESULTS_DIR/export/ocdnet_model.engine \
inference.input_folder=$DATA_DIR/test/img \
inference.results_dir=$RESULTS_DIR/inference