Object Detection using TAO DetectNet_v2

Transfer learning is the process of transferring learned features from one application to another. It is a commonly used training technique where you use a model trained on one task and re-train to use it on a different task.

Train Adapt Optimize (TAO) Toolkit is a simple and easy-to-use Python based AI toolkit for taking purpose-built AI models and customizing them with users' own data.

Learning Objectives

In this notebook, you will learn how to leverage the simplicity and convenience of TAO to:

• Take a pretrained resnet18 model and train a ResNet-18 DetectNet_v2 model on the KITTI dataset
• Prune the trained detectnet_v2 model
• Retrain the pruned model to recover lost accuracy
• Export the pruned model
• Quantize the pruned model using QAT
• Run Inference on the trained model
• Export the pruned, quantized and retrained model to a .etlt file for deployment to DeepStream
• Run inference on the exported. etlt model to verify deployment using TensorRT

Table of Contents

This notebook shows an example usecase of Object Detection using DetectNet_v2 in the Train Adapt Optimize (TAO) Toolkit.

0. Set up env variables and map drives
1. Install the TAO Launcher
2. Prepare dataset and pre-trained model
   1. Download the dataset
   2. Verify downloaded dataset
   3. Prepare tfrecords from kitti format dataset
   4. Download pre-trained model
3. Provide training specification
4. Run TAO training
5. Evaluate trained models
6. Prune trained models
7. Retrain pruned models
8. Evaluate retrained model
9. Visualize inferences
10. Model Export
    1. Int8 Optimization
    2. Generate TensorRT engine
11. Verify Deployed Model
    1. Inference using TensorRT engine
12. QAT workflow
    1. Convert pruned model to QAT and retrain
    2. Evaluate QAT converted model
    3. Export QAT trained model to int8
    4. Evaluate a QAT trained model using the exported TensorRT engine
    5. Inference using QAT engine

0. Set up env variables and map drives

When using the purpose-built pretrained models from NGC, please make sure to set the $KEY environment variable to the key as mentioned in the model overview. Failing to do so, can lead to errors when trying to load them as pretrained models.

The following notebook requires the user to set an env variable called the $LOCAL_PROJECT_DIR as the path to the users workspace. Please note that the dataset to run this notebook is expected to reside in the $LOCAL_PROJECT_DIR/data, while the TAO experiment generated collaterals will be output to $LOCAL_PROJECT_DIR/detectnet_v2. More information on how to set up the dataset and the supported steps in the TAO workflow are provided in the subsequent cells.

Note: Please make sure to remove any stray artifacts/files from the $USER_EXPERIMENT_DIR or $DATA_DOWNLOAD_DIR paths as mentioned below, that may have been generated from previous experiments. Having checkpoint files etc may interfere with creating a training graph for a new experiment.

Note: This notebook currently is by default set up to run training using 1 GPU. To use more GPU's please update the env variable $NUM_GPUS accordingly

The cell below maps the project directory on your local host to a workspace directory in the TAO docker instance, so that the data and the results are mapped from in and out of the docker. For more information please refer to the launcher instance in the user guide.

When running this cell on AWS, update the drive_map entry with the dictionary defined below, so that you don't have permission issues when writing data into folders created by the TAO docker.

drive_map = {
    "Mounts": [
            # Mapping the data directory
            {
                "source": os.environ["LOCAL_PROJECT_DIR"],
                "destination": "/workspace/tao-experiments"
            },
            # Mapping the specs directory.
            {
                "source": os.environ["LOCAL_SPECS_DIR"],
                "destination": os.environ["SPECS_DIR"]
            },
        ],
    "DockerOptions": {
        "user": "{}:{}".format(os.getuid(), os.getgid())
    }
}

1. Install the TAO launcher

The TAO launcher is a python package distributed as a python wheel listed in the nvidia-pyindex python index. You may install the launcher by executing the following cell.

Please note that TAO Toolkit recommends users to run the TAO launcher in a virtual env with python 3.6.9. You may follow the instruction in this page to set up a python virtual env using the virtualenv and virtualenvwrapper packages. Once you have setup virtualenvwrapper, please set the version of python to be used in the virtual env by using the VIRTUALENVWRAPPER_PYTHON variable. You may do so by running

export VIRTUALENVWRAPPER_PYTHON=/path/to/bin/python3.x

where x >= 6 and <= 8

We recommend performing this step first and then launching the notebook from the virtual environment. In addition to installing TAO python package, please make sure of the following software requirements:

• python >=3.6.9 < 3.8.x
• docker-ce > 19.03.5
• docker-API 1.40
• nvidia-container-toolkit > 1.3.0-1
• nvidia-container-runtime > 3.4.0-1
• nvidia-docker2 > 2.5.0-1
• nvidia-driver > 455+

Once you have installed the pre-requisites, please log in to the docker registry nvcr.io by following the command below

docker login nvcr.io

You will be triggered to enter a username and password. The username is $oauthtoken and the password is the API key generated from ngc.nvidia.com. Please follow the instructions in the NGC setup guide to generate your own API key.

2. Prepare dataset and pre-trained model

We will be using the kitti object detection dataset for this example. To find more details, please visit http://www.cvlibs.net/datasets/kitti/eval_object.php?obj_benchmark=2d. Please download both, the left color images of the object dataset from here and, the training labels for the object dataset from here, and place the zip files in $LOCAL_DATA_DIR

The data will then be extracted to have

• training images in $LOCAL_DATA_DIR/training/image_2
• training labels in $LOCAL_DATA_DIR/training/label_2
• testing images in $LOCAL_DATA_DIR/testing/image_2

You may use this notebook with your own dataset as well. To use this example with your own dataset, please follow the same directory structure as mentioned below.

Note: There are no labels for the testing images, therefore we use it just to visualize inferences for the trained model.

A. Download the dataset

Once you have gotten the download links in your email, please populate them in place of the KITTI_IMAGES_DOWNLOAD_URL and the KITTI_LABELS_DOWNLOAD_URL. This next cell, will download the data and place in $LOCAL_DATA_DIR

B. Verify downloaded dataset

C. Prepare tf records from kitti format dataset

• Update the tfrecords spec file to take in your kitti format dataset
• Create the tfrecords using the detectnet_v2 dataset_convert

Note: TfRecords only need to be generated once.

D. Download pre-trained model

Download the correct pretrained model from the NGC model registry for your experiment. Please note that for DetectNet_v2, the input is expected to be 0-1 normalized with input channels in RGB order. Therefore, for optimum results please download model templates from nvidia/tao/pretrained_detectnet_v2. The templates are now organized as version strings. For example, to download a resnet18 model suitable for detectnet please resolve to the ngc object shown as nvidia/tao/pretrained_detectnet_v2:resnet18.

All other models are in BGR order expect input preprocessing with mean subtraction and input channels. Using them as pretrained weights may result in suboptimal performance.

You may also use this notebook with the following purpose-built pretrained models

• PeopleNet
• TrafficCamNet
• DashCamNet
• FaceDetect-IR

3. Provide training specification

• Tfrecords for the train datasets
  • To use the newly generated tfrecords, update the dataset_config parameter in the spec file at $SPECS_DIR/detectnet_v2_train_resnet18_kitti.txt
  • Update the fold number to use for evaluation. In case of random data split, please use fold 0 only
  • For sequence-wise split, you may use any fold generated from the dataset convert tool
• Pre-trained models
• Augmentation parameters for on the fly data augmentation
• Other training (hyper-)parameters such as batch size, number of epochs, learning rate etc.

4. Run TAO training

• Provide the sample spec file and the output directory location for models

Note: The training may take hours to complete. Also, the remaining notebook, assumes that the training was done in single-GPU mode. When run in multi-GPU mode, please expect to update the pruning and inference steps with new pruning thresholds and updated parameters in the clusterfile.json accordingly for optimum performance.

Detectnet_v2 now supports restart from checkpoint. In case the training job is killed prematurely, you may resume training from the closest checkpoint by simply re-running the same command line. Please do make sure to use the same number of GPUs when restarting the training.

When running the training with NUM_GPUs>1, you may need to modify the batch_size_per_gpu and learning_rate to get similar mAP as a 1GPU training run. In most cases, scaling down the batch-size by a factor of NUM_GPU's or scaling up the learning rate by a factor of NUM_GPU's would be a good place to start.

5. Evaluate the trained model

6. Prune the trained model

• Specify pre-trained model
• Equalization criterion (Applicable for resnets and mobilenets)
• Threshold for pruning.
• A key to save and load the model
• Output directory to store the model

Usually, you just need to adjust -pth (threshold) for accuracy and model size trade off. Higher pth gives you smaller model (and thus higher inference speed) but worse accuracy. The threshold to use is dependent on the dataset. A pth value 5.2e-6 is just a start point. If the retrain accuracy is good, you can increase this value to get smaller models. Otherwise, lower this value to get better accuracy.

For some internal studies, we have noticed that a pth value of 0.01 is a good starting point for detectnet_v2 models.

7. Retrain the pruned model

• Model needs to be re-trained to bring back accuracy after pruning
• Specify re-training specification with pretrained weights as pruned model.

Note: For retraining, please set the load_graph option to true in the model_config to load the pruned model graph. Also, if after retraining, the model shows some decrease in mAP, it could be that the originally trained model was pruned a little too much. Please try reducing the pruning threshold (thereby reducing the pruning ratio) and use the new model to retrain.

Note: DetectNet_v2 now supports Quantization Aware Training, to help with optmizing the model. By default, the training in the cell below doesn't run the model with QAT enabled. For information on training a model with QAT, please refer to the cells under section 11

8. Evaluate the retrained model

This section evaluates the pruned and retrained model, using the evaluate command.

9. Visualize inferences

In this section, we run the inference tool to generate inferences on the trained models. To render bboxes from more classes, please edit the spec file detectnet_v2_inference_kitti_tlt.txt to include all the classes you would like to visualize and edit the rest of the file accordingly.

The inference tool produces two outputs.

1. Overlain images in $USER_EXPERIMENT_DIR/tlt_infer_testing/images_annotated
2. Frame by frame bbox labels in kitti format located in $USER_EXPERIMENT_DIR/tlt_infer_testing/labels

Note: To run inferences for a single image, simply replace the path to the -i flag in inference command with the path to the image.

10. Model Export

A. Int8 Optimization

DetectNet_v2 model supports int8 inference mode in TensorRT. In order to use int8 mode, we must calibrate the model to run 8-bit inferences -

• Generate calibration tensorfile from the training data using detectnet_v2 calibration_tensorfile
• Use tao export to generate int8 calibration table.

Note: For this example, we generate a calibration tensorfile containing 10 batches of training data. Ideally, it is best to use at least 10-20% of the training data to do so. The more data provided during calibration, the closer int8 inferences are to fp32 inferences.

Note: If the model was trained with QAT nodes available, please refrain from using the post training int8 optimization as mentioned below. Please export the model in int8 mode (using the arg --data_type int8) with just the path to the calibration cache file (using the argument --cal_cache_file)

B. Generate TensorRT engine

Verify engine generation using the tao-converter utility included with the docker.

The tao-converter produces optimized tensorrt engines for the platform that it resides on. Therefore, to get maximum performance, please instantiate this docker and execute the tao-converter command, with the exported .etlt file and calibration cache (for int8 mode) on your target device. The tao-converter utility included in this docker only works for x86 devices, with discrete NVIDIA GPU's.

For the jetson devices, please download the tao-converter for jetson from the dev zone link here.

If you choose to integrate your model into deepstream directly, you may do so by simply copying the exported .etlt file along with the calibration cache to the target device and updating the spec file that configures the gst-nvinfer element to point to this newly exported model. Usually this file is called config_infer_primary.txt for detection models and config_infer_secondary_*.txt for classification models.

11. Verify Deployed Model

Verify the exported model by visualizing inferences on TensorRT. In addition to running inference on a .tlt model in step 9, the inference tool is also capable of consuming the converted TensorRT engine from step 10.B.

*If after int-8 calibration the accuracy of the int-8 inferences seem to degrade, it could be because the there wasn't enough data in the calibration tensorfile used to calibrate thee model or, the training data is not entirely representative of your test images, and the calibration maybe incorrect. Therefore, you may either regenerate the calibration tensorfile with more batches of the training data and recalibrate the model, or calibrate the model on a few images from the test set. This may be done using --cal_image_dir flag in the export tool. For more information, please follow the instructions in the USER GUIDE.

A. Inference using TensorRT engine

11. QAT workflow

This section delves into the newly enabled Quantization Aware Training feature with DetectNet_v2. The workflow defined below converts a pruned model from section 5 to enable QAT and retrain this model to while accounting the noise introduced due to quantization in the forward pass.

A. Convert pruned model to QAT and retrain

All detectnet models, unpruned and pruned models can be converted to QAT models by setting the enable_qat parameter in the training_config component of the spec file to true.

B. Evaluate QAT converted model

This section evaluates a QAT enabled pruned retrained model. The mAP of this model should be comparable to that of the pruned retrained model without QAT. However, due to quantization, it is possible sometimes to see a drop in the mAP value for certain datasets.

C. Export QAT trained model to int8

Export a QAT trained model to TensorRT parsable model. This command generates an .etlt file from the trained model and the serializes corresponding int8 scales as a TRT readable calibration cache file.

D. Evaluate a QAT trained model using the exported TensorRT engine

This section evaluates a QAT enabled pruned retrained model using the TensorRT int8 engine that was exported in Section C. Please note that there maybe a slight difference (~0.1-0.5%) in the mAP from Section B, oweing to some differences in the implementation of quantization in TensorRT.

Note: The TensorRT evaluator might be slightly slower than the TAO evaluator here, because the evaluation dataloader is pinned to the CPU to avoid any clashes between TensorRT and TAO instances in the GPU. Please note that this tool was not intended and has not been developed for profiling the model. It is just a means to qualitatively analyse the model.

Please use native TensorRT or DeepStream for the most optimized inferences.

E. Inference using QAT engine

Run inference and visualize detections on test images, using the exported TensorRT engine from Section C.