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.
Access the satellite detection model
To replicate these results, you can clone the GitHub repository and follow along with the included Jupyter notebook.
Clone the following repo:
git clone git@github.com:aireveries/rareplanes-tlt.git ~/Code/rareplanes-tlt
Create a conda
environment:
conda create -f env.yaml
Activate the model:
source activate rareplanes-tlt
Start Jupyter:
jupyter notebook
Learning objectives
- 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 the NVIDIA Container Toolkit / nvidia-docker2. For more information, see the NVIDIA Container Toolkit Installation Guide.
- 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.
Downloading the datasets
For more information about the contents of the RarePlanes dataset, see RarePlanes Public User Guide.
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:
!mkdir -p data/real/tarballs/{train,test} !mkdir -p data/synthetic
Now use this function to download the datasets from Amazon S3, extract them, and verify:
def download(s3_path, out_folder, out_file_count): rel_file_path = Path('data') / Path(s3_path.replace('s3://rareplanes-public/', '')) rel_folder = rel_file_path.parent / out_folder num_files = !ls $rel_folder | wc -l try: if int(num_files[0]) == out_file_count: print(f'{s3_path} already downloaded and extracted') else: raise Exception except: if not rel_file_path.exists(): print('Starting download') !aws s3 cp $s3_path $rel_file_path; else: print(f'{s3_path} already downloaded') print('Extracting...') !cd {rel_folder.parent}; pv {rel_file_path.name} | tar xz; print('Removing compressed file.') !rm $rel_file_path
Then download the dataset:
download('s3://rareplanes-public/real/tarballs/metadata_annotations.tar.gz', 'metadata_annotations', 9) download('s3://rareplanes-public/real/tarballs/train/RarePlanes_train_PS-RGB_tiled.tar.gz', 'PS-RGB_tiled', 11630) download('s3://rareplanes-public/real/tarballs/test/RarePlanes_test_PS-RGB_tiled.tar.gz', 'PS-RGB_tiled', 5420) !aws s3 cp --recursive s3://rareplanes-public/synthetic/ data/synthetic
Converting from COCO to KITTI format
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.
{ "Mounts": [ { "source": "/home/patrick.rodriguez/Code/rareplanes-tlt", "destination": "/workspace/tlt-experiments" } ], "Envs": [ { "variable": "CUDA_VISIBLE_DEVICES", "value": "0" } ], "DockerOptions": { "shm_size": "16G", "ulimits": { "memlock": -1, "stack": 67108864 }, "user": "1001:1001" } }
Processing datasets into TFRecords
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
The model is now located at the following path:
./tlt_pretrained_detectnet_v2_vresnet18/resnet18.hdf5
Run a benchmark experiment using real data
The following command starts training and logs results to a file that you can tail:
!tlt detectnet_v2 train --key tlt --gpu_index 0 -e /workspace/tlt-experiments/specs/detectnet_v2_train_resnet18_kitti_real.txt -r /workspace/tlt-experiments/detectnet_v2_outputs/resnet18_real_amp16 -n resnet18_real_amp16 --use_amp > out_resnet18_real_amp16.log
Follow along with the following command:
tail -f ./out_resnet18_real_amp16.log
After training is complete, you can use the functions defined in the notebook to get relevant statistics on your model:
get_model_param_counts('./out_resnet18_real_amp16.log') best_epoch = get_best_epoch('./out_resnet18_real_amp16.log') best_epoch
You get something like the following output:
Total params: 11,197,893 Trainable params: 11,188,165 Non-trainable params: 9,728 Best epoch and map50 metric: (79, 94.2296)
To reevaluate your trained model on your test set or other dataset, run the following:
!tlt detectnet_v2 evaluate --gpu_index 0 -e /workspace/tlt-experiments/specs/detectnet_v2_evaluate_real.txt -m /workspace/tlt-experiments/{best_checkpoint} -k tlt
The output should look something like this:
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.
Running an experiment with synthetic data
!tlt detectnet_v2 train --key tlt --gpu_index 0 -e /workspace/tlt-experiments/specs/detectnet_v2_train_resnet18_kitti_synth.txt -r /workspace/tlt-experiments/detectnet_v2_outputs/resnet18_synth_amp16 -n resnet18_synth_amp16 --use_amp > out_resnet18_synth_amp16.log
You can see the results for each epoch by running: !cat out_resnet18_synth_amp16.log | grep -i aircraft
Example output:
aircraft 58.1444 aircraft 65.1423 aircraft 64.3203 aircraft 68.1934 aircraft 71.5754 aircraft 68.5568
Fine-tuning the synthetic-trained model with real data
Now, fine-tune your best-performing synthetic-data-trained model with 10% of the real data. To do so, you must first create the 10% split.
%run ./create_train_split.py convert_split('kitti_real_train_10')
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.
!tlt detectnet_v2 train --key tlt --gpu_index 0 -e /workspace/tlt-experiments/specs/detectnet_v2_train_resnet18_kitti_synth_finetune_10_replaced.txt -r /workspace/tlt-experiments/detectnet_v2_outputs/resnet18_synth_finetune_10_amp16 -n resnet18_synth_finetune_10_amp16 --use_amp > out_resnet18_synth_finetune_10_amp16.log
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:
!tlt detectnet_v2 evaluate --gpu_index 0 -e /workspace/tlt-experiments/specs/detectnet_v2_evaluate_real.txt -m /workspace/tlt-experiments/{best_checkpoint} -k tlt
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:00Data 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.!mkdir -p detectnet_v2_outputs/pruned !tlt detectnet_v2 prune -m /workspace/tlt-experiments/{best_checkpoint} -o /workspace/tlt-experiments/detectnet_v2_outputs/pruned/pruned-model.tlt -eq union -pth 0.5 -k tltYou can now evaluate the pruned model:
!tlt detectnet_v2 evaluate --gpu_index 0 -e /workspace/tlt-experiments/specs/detectnet_v2_evaluate_real.txt -m /workspace/tlt-experiments/detectnet_v2_outputs/pruned/pruned-model.tlt -k tlt > out_pruned.txtNow you can see how many parameters remain:
get_model_param_counts('./out_pruned.txt')You should see something like the following outputs:
Total params: 3,372,973 Trainable params: 3,366,573 Non-trainable params: 6,400This 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:
!cat out_pruned.txt | grep -i aircraft aircraft 68.8865Luckily, 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)You can now retrain the pruned model:
!tlt detectnet_v2 train --key tlt --gpu_index 0 -e /workspace/tlt-experiments/specs/detectnet_v2_train_resnet18_kitti_synth_finetune_10_pruned_retrain_replaced.txt -r /workspace/tlt-experiments/detectnet_v2_outputs/resnet18_synth_finetune_10_pruned_retrain_amp16 -n resnet18_synth_finetune_10_pruned_retrain_amp16 --use_amp > out_resnet18_synth_finetune_10_pruned_retrain_amp16.logOn 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:01Quantizing 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)Run the QAT training:
!tlt detectnet_v2 train --key tlt --gpu_index 0 -e /workspace/tlt-experiments/specs/detectnet_v2_train_resnet18_kitti_synth_finetune_10_pruned_retrain_qat_replaced.txt -r /workspace/tlt-experiments/detectnet_v2_outputs/resnet18_synth_finetune_10_pruned_retrain_qat_amp16 -n resnet18_synth_finetune_10_pruned_retrain_qat_amp16 --use_amp > out_resnet18_synth_finetune_10_pruned_retrain_qat_amp16.logUse the TLT export tool to export to INT8 quantized TensorRT format:
!tlt detectnet_v2 export -m /workspace/tlt-experiments/{best_checkpoint} -o /workspace/tlt-experiments/detectnet_v2_outputs/qat/resnet18_detector_qat.etlt -k tlt --data_type int8 --batch_size 64 --max_batch_size 64 --engine_file /workspace/tlt-experiments/detectnet_v2_outputs/qat/resnet18_detector_qat.trt.int8 --cal_cache_file /workspace/tlt-experiments/detectnet_v2_outputs/qat/calibration_qat.bin --verboseAt this point, you can now evaluate your quantized model using TensorRT:
!tlt detectnet_v2 evaluate -e /workspace/tlt-experiments/specs/detectnet_v2_train_resnet18_kitti_synth_finetune_10_pruned_retrain_qat_replaced.txt -m /workspace/tlt-experiments/detectnet_v2_outputs/qat/resnet18_detector_qat.trt.int8 -f tensorrtLooking at the output:
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:01Conclusion
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.
Go to AI.Reverie, download the synthetic training data for your project, and start training with TLT.