CloudSoftwareHardware
Software Stack
DepthAI
OAK Apps
AI Models
Troubleshooting

ON THIS PAGE

  • NVIDIA TAO
  • Overview
  • Usage

NVIDIA TAO

Overview

NVIDIA TAO (Train, Adapt, and Optimize) is a powerful AI toolkit designed to simplify and accelerate the development of deep learning models for computer vision, speech, and natural language processing. Built on NVIDIA's cutting-edge AI frameworks, TAO enables developers to fine-tune pre-trained models with minimal coding, reducing the need for large datasets and extensive machine-learning expertise. By leveraging transfer learning and model pruning, TAO streamlines the AI development pipeline, making it accessible for enterprises looking to deploy high-performance AI applications efficiently. You can learn more about the toolkit in their official documentation.In addition to the toolkit they also provide a NCG Catalog which amongst things like collections, containers, and helm charts also offers a vast Model Zoo for many different use cases. You can explore all the models here.

Usage

TAO Toolkit allows you to train a custom model. For detailed guidance, refer to the official documentation and explore various hands-on tutorials. Once training is complete, you'll have a .tlt model file, which needs to be converted to the ONNX format. Below is an example of how to export a re-identification model:
Command Line
1tao model re_identification export \
2 -e $USER_EXPERIMENT_DIR/media/experiment_market1501_resnet.yaml \
3    results_dir=$USER_EXPERIMENT_DIR \
4 encryption_key=$KEY \
5    export.checkpoint=$USER_EXPERIMENT_DIR/resnet50_market1501_aicity156.tlt \
6 export.onnx_file=$USER_EXPERIMENT_DIR/resnet50_market1501_aicity156.onnx
After exporting, you should have an ONNX file. From here, follow the steps outlined in the RVC Conversion section. Alternatively, you can refer to our ONNX conversion tutorial for a detailed walkthrough.If you prefer not to train your own model, you can use a pre-trained model from the TAO Model Zoo. Many of these models are already available in ONNX format, allowing you to skip the conversion step and proceed directly with preparing the model for deployment on an RVC device.