The long, cumbersome slog of data procurement has been slowing down innovation in AI, especially in computer vision, which relies on labeled images and video for training. But now you can jumpstart your machine learning process by quickly generating synthetic data using AI.Reverie. With the AI.Reverie synthetic data platform, you can create the exact training … Continued
The long, cumbersome slog of data procurement has been slowing down innovation in AI, especially in computer vision, which relies on labeled images and video for training. But now you can jumpstart your machine learning process by quickly generating synthetic data using AI.Reverie.
With the AI.Reverie synthetic data platform, you can create the exact training data that you need in a fraction of the time it would take to find and label the right real photography. In AI.Reverie’s photorealistic 3D environments, you can generate data for all possible scenarios, including hard to reach places, unusual environmental conditions, and rare or unique events.
Training data generation includes labels. Choose the needed types, such as 2D or 3D bounding boxes, depth masks, and so on. After you test your model, you can return to the platform to quickly generate additional data to improve accuracy. Test and repeat in quick, iterative cycles.
We wanted to test performance of AI.Reverie synthetic data in NVIDIA Transfer Learning Toolkit 3.0. Originally, we set out to replicate the results in the research paper RarePlanes: Synthetic Data Takes Flight, which used synthetic imagery to create object detection models. We discovered new tools in TLT that made it possible to create more lightweight models that were as accurate as, but much faster than, those featured in the original paper.
In this post, we show you how we used the TLT quantized-aware training and model pruning to accomplish this, and how to replicate the results yourself. We show you how to create an airplane detector, but you should be able to fine-tune the model for various satellite detection scenarios of your own.
Figure 1. A synthetic image featuring annotations that denote aircraft type, wing shape, and other distinguishing features.
Access the satellite detection model
To replicate these results, you can clone the GitHub repository and follow along with the included Jupyter notebook.
Generate synthetic data using the AI.Reverie platform and use it with NVIDIA TLT.
Train highly accurate models using synthetic data.
Optimize a model for inference using the TLT.
Prerequisites
We tested the code with Python 3.8.8, using Anaconda 4.9.2 to manage dependencies and the virtual environment. The code may work with different versions of Python and other virtual environment solutions, but we haven’t tested those configurations. We used Ubuntu 18.04.5 LTS and NVIDIA driver 460.32.03 and CUDA Version 11.2. TLT requires driver 455.xx or later.
Set up NGC to be able to download NVIDIA Docker containers. Follow steps 4 and 5 in the TLT User Guide. For more information about the NGC CLI tool, see CLI Install.
Have available at least 250 GB hard disk space to store dataset and model weights.
For this tutorial, you need only download a subset of the data. The following code example is meant to be executed from within the Jupyter notebook. First, create the folders:
TLT uses the KITTI format for object detection model training. RarePlanes is in the COCO format, so you must run a conversion script from within the Jupyter notebook. This converts the real train/test and synthetic train/test datasets.
%run convert_coco_to_kitti.py
There should now be a folder for each dataset split inside of data/kitti that contains the KITTI formatted annotation text files and symlinks to the original images.
Setting up TLT mounts
The notebook has a script to generate a ~/.tlt_mounts.json file. For more information about the various settings, see Running the launcher.
You must turn the KITTI labels into the TFRecord format used by TLT. The convert_split function in the notebook helps you bulk convert all the datasets:
def convert_split(name):
!tlt detectnet_v2 dataset_convert --gpu_index 0
-d /workspace/tlt-experiments/specs/detectnet_v2_tfrecords_{name}.txt
-o /workspace/tlt-experiments/data/tfrecords/{name}/{name}
You can then run the conversions:
convert_split('kitti_real_train')
convert_split('kitti_real_test')
convert_split('kitti_synthetic_train')
convert_split('kitti_synthetic_test')
Download the ResNet18 convolutional backbone
Using your NGC account and command-line tool, you can now download the model:
Download the ResNet18 convolutional backbone
Using your NGC account and command-line tool, you can now download the model:
!ngc registry model download-version nvidia/tlt_pretrained_detectnet_v2:resnet18
Validation cost: 0.001133
Mean average_precision (in %): 94.2563
class name average precision (in %)
------------ --------------------------
aircraft 94.2563
Median Inference Time: 0.003877
2021-04-06 05:47:00,323 [INFO] __main__: Evaluation complete.
Time taken to run __main__:main: 0:00:27.031500.
2021-04-06 05:47:02,466 [INFO] tlt.components.docker_handler.docker_handler: Stopping container.
You then use this function to replace the checkpoint in your template spec with the best performing model from the synthetic-only training.
with open('./specs/detectnet_v2_train_resnet18_kitti_synth_finetune_10.txt', 'r') as f_in:
with open('./specs/detectnet_v2_train_resnet18_kitti_synth_finetune_10_replaced.txt', 'w') as f_out:
out = f_in.read().replace('REPLACE', best_checkpoint)
f_out.write(out)
You can now begin a TLT training. Start your fine-tuning with the best-performing epoch of the model trained on synthetic data alone, in the previous section.
After training has completed, you should see a best epoch of between 91-93% mAP50, which gets you close to the real-only model performance with only 10% of the real data.
In the notebook, there’s a command to evaluate the best performing model checkpoint on the test set:
You should see something like the following output:
2021-04-06 18:05:28,342 [INFO] iva.detectnet_v2.evaluation.evaluation: step 330 / 339, 0.05s/step
Matching predictions to ground truth, class 1/1.: 100%|█| 14719/14719 [00:00
Figure 2. Training on synthetic + 10% real data nearly matches the results of training on 100% of the real data.
Data enhancement is fine-tuning a model training on AI.Reverie’s synthetic data with just 10% of the original, real dataset. As you can see, this technique produces a model as accurate as one trained on real data alone. That represents roughly 90% cost savings on real, labeled data and saves you from having to endure a long hand-labeling and QA process.
Pruning the model
Having trained a well-performing model, you can now decrease the number of weights to cut down on file size and inference time. TLT includes an easy-to-use pruning tool.
The one argument to play with is -pth, which sets the threshold for neurons to prune. The higher you set this, the more parameters are pruned, but after a certain point your accuracy metric may drop too low. We found that a value of 0.5 worked for these experiments, but you may find different results on other datasets.
You should see something like the following outputs:
Total params: 3,372,973
Trainable params: 3,366,573
Non-trainable params: 6,400
This is 70% smaller than the original model, which had 11.2 million parameters! Of course, you’ve lost performance by dropping so many parameters, which you can verify:
Luckily, you can recover almost all the performance by retraining the pruned model.
Retraining the models
As before, there is a template spec to run this experiment that only requires you to fill in the location of the pruned model:
with open('./specs/detectnet_v2_train_resnet18_kitti_synth_finetune_10_pruned_retrain.txt', 'r') as f_in:
with open('./specs/detectnet_v2_train_resnet18_kitti_synth_finetune_10_pruned_retrain_replaced.txt', 'w') as f_out:
out = f_in.read().replace('REPLACE', 'detectnet_v2_outputs/pruned/pruned-model.tlt')
f_out.write(out)
On a run of this experiment, the best performing epoch achieved 91.925 mAP50, which is about the same as the original nonpruned experiment.
2021-04-06 19:33:39,360 [INFO] iva.detectnet_v2.evaluation.evaluation: step 330 / 339, 0.05s/step
Matching predictions to ground truth, class 1/1.: 100%|█| 17403/17403 [00:01
Quantizing the models
The final step in this process is quantizing the pruned model so that you can achieve much higher levels of inference speed with TensorRT. We have a quantization aware training (QAT) spec template available:
with open('./specs/detectnet_v2_train_resnet18_kitti_synth_finetune_10_pruned_retrain_qat.txt', 'r') as f_in:
with open('./specs/detectnet_v2_train_resnet18_kitti_synth_finetune_10_pruned_retrain_qat_replaced.txt', 'w') as f_out:
out = f_in.read().replace('REPLACE', 'detectnet_v2_outputs/pruned/pruned-model.tlt')
f_out.write(out)
2021-04-06 23:08:28,471 [INFO] iva.detectnet_v2.evaluation.tensorrt_evaluator: step 330 / 339, 0.33s/step
Matching predictions to ground truth, class 1/1.: 100%|█| 21973/21973 [00:01
Conclusion
We were impressed by these results. AI.Reverie’s synthetic data platform, with just 10% of the real dataset, enabled us to achieve the same performance as we did when training on the full real dataset. That represents a cost savings of roughly 90%, not to mention the time saved on procurement. It now takes days, not months, to generate the needed synthetic data.
TLT also produced a 25.2x reduction in parameter count, a 33.6x reduction in file size, a 174.7x increase in performance (QPS), while retaining 95% of the original performance. TLT’s capabilities were particularly valuable for pruning and quantizing.