TAO Toolkit Quick Start
The NVIDIA TAO Toolkit, built on TensorFlow and PyTorch, simplifies and accelerates the model training process by abstracting away the complexity of AI models and the deep learning framework. You can use the power of transfer learning to fine-tune NVIDIA pretrained models with your own data and optimize the model for inference throughput — all without the need for AI expertise or large training datasets.
TAO Toolkit 5.3
TAO Toolkit 5.3 version packages containers, models, Jupyter notebooks, start-up script and Helm chart for K8s deployment. Here are list of assets as part of TAO 5.3
Containers
The containers for the TAO Toolkit 5.3 are hosted on NGC under this instance.
| Container Name | Description | TAG | GitHub Repo |
|---|---|---|---|
| TAO TensorFlow 1 | TensorFlow 1.15.x container for training DNNs | nvcr.io/nvidia/tao/tao-toolkit:5.0.0-tf1.15.5 | https://github.com/NVIDIA/tao_tensorflow1_backend |
| TAO TensorFlow 2 | TensorFlow 2.11.x container for training DNNs | nvcr.io/nvidia/tao/tao-toolkit:5.0.0-tf2.11.0 | https://github.com/NVIDIA/tao_tensorflow2_backend |
| TAO PyTorch | PyTorch container for training DNNs | nvcr.io/nvidia/tao/tao-toolkit:5.3.0-pyt | https://github.com/NVIDIA/tao_pytorch_backend |
| TAO Deploy | TensorRT container for optimization | nvcr.io/nvidia/tao/tao-toolkit:5.3.0-deploy | https://github.com/NVIDIA/tao_deploy |
| TAO Data Service | Container for AI-assisted annotation and few other data services | nvcr.io/nvidia/tao/tao-toolkit:5.3.0-data-services | https://github.com/NVIDIA/tao_dataset_suite |
| TAO Services | Container for TAO services | nvcr.io/nvidia/tao/tao-toolkit:5.3.0-api | https://github.com/NVIDIA/tao_front_end_services |
Resources
| Container Name | Description | Location |
|---|---|---|
| TAO Toolkit Getting Started | Resource to help get started with TAO Toolkit. | nvidia/tao/tao-getting-started |
Helm Chart
The TAO Toolkit Services Helm chart is hosted on NGC at nvidia/tao/tao-toolkit-api
Models
The following models are released as part of the TAO Toolkit.
| Model name | Description | **NGC Instance |
|---|---|---|
| EmotionNet | Network to classify emotions from face. | nvidia/tao/emotionnet |
| TrafficCamNet | 4 class object detection network to detect cars in an image. | nvidia/tao/trafficcamnet |
| Optical Character Detection | Network to detect characters in an image. | nvidia/tao/ocdnet |
| CitySemSegformer | Model to segment objects in an image. | nvidia/tao/citysemsegformer |
| DashCamNet | 4 class object detection network to detect cars in an image. | nvidia/tao/dashcamnet |
| PeopleSemSegnet | Semantic segmentation of persons in an image. | nvidia/tao/peoplesemsegnet |
| PeopleNet | 3 class object detection network to detect people in an image. | nvidia/tao/peoplenet |
| CenterPose - Isaac ROS | Single-stage, keypoint-based method for category-level object pose estimation. | nvidia/tao/centerpose_ros |
| Re-Identification | Re-Identification network to generate embeddings for identifying persons in different scenes. | nvidia/tao/reidentificationnet--discoverable |
| Deformable DETR | Model trained on COCO2017 to detect and classify objects. | nvidia/tao/pretrained_deformable_detr_coco |
| VehicleTypeNet | Resnet18 model to classify a car crop into 1 out 6 car types. | nvidia/tao/vehicletypenet |
| Action Recognition Net | 5 class action recognition network to recognize what people do in an image. | nvidia/tao/actionrecognitionnet |
| LPDNet | Object Detection network to detect license plates in an image of a car. | nvidia/tao/lpdnet |
| GestureNet | Classify gestures from hand crop images. | nvidia/tao/gesturenet |
| BodyPoseNet | Detect body pose from an image. | nvidia/tao/bodyposenet |
| License Plate Recognition | Model to recognize characters from the image crop of a License Plate. | nvidia/tao/lprnet |
| FaceDetectIR | 1 class object detection network to detect faces in an image. | nvidia/tao/facedetectir |
| Re-Identification | Re-Identification network to generate embeddings for identifying persons in different scenes. | nvidia/tao/reidentificationnet |
| Visual ChangeNet Segmentation - (Research-only) | Visual ChangeNet-Segmentation (Research-only) | nvidia/tao/visual_changenet_segmentation_levircd |
| Visual ChangeNet Segmentation - (Commercial) | Visual ChangeNet-Segmentation (Commercial) | nvidia/tao/visual_changenet_segmentation_landsatscd |
| DINO | Model trained on COCO2017 to detect and classify objects. | nvidia/tao/pretrained_dino_coco |
| PointPillarNet | Model to detect one or more objects from a LIDAR point cloud file and return 3D bounding boxes. | nvidia/tao/pointpillarnet |
| PeopleSemSegformer | Model to segment persons in an image. | nvidia/tao/peoplesemsegformer |
| Optical Character Recognition | Model to recognise characters from a preceding OCDNet model. | nvidia/tao/ocrnet |
| HeartRateNet | Estimate heart-rate non-invasively from RGB facial videos. | nvidia/tao/heartratenet |
| Retail Object Detection | EfficientDet based object detection network to detect retail objects on a checkout counter. | nvidia/tao/retail_object_detection |
| FaceDetect | Detect faces from an image. | nvidia/tao/facenet |
| Optical Inspection | Model to detect defects in soldered components on a Printed Circuit Board. | nvidia/tao/optical_inspection |
| PeopleNet - AMR | 3 class object detection network to detect people in an image. | nvidia/tao/peoplenet_amr |
| Mask Auto Label | Pretrained model to generate semantic segmentation labels. | nvidia/tao/mask_auto_label |
| Retail Object Recognition | Embedding generator model to recognize objects on a checkout counter. | nvidia/tao/retail_object_recognition |
| PeopleNet Transformer | 3 class object detection network to detect people in an image. | nvidia/tao/peoplenet_transformer |
| Visual ChangeNet Classification | Visual ChangeNet - Classification Models | nvidia/tao/visual_changenet_classification |
| PeopleSegNet | 1 class instance segmentation network to detect and segment instances of people in an image. | nvidia/tao/peoplesegnet |
| VehicleMakeNet | Resnet18 model to classify a car crop into 1 out 20 car brands. | nvidia/tao/vehiclemakenet |
| Facial Landmarks Estimation | Detect fiducial keypoints from an image of a face. | nvidia/tao/fpenet |
| Pose Classification | Pose classification network to classify poses of people from their skeletons. | nvidia/tao/poseclassificationnet |
| BodyPose3DNet | 3D human pose estimation network to predict 34 keypoints in 3D of a person in an image. | nvidia/tao/bodypose3dnet |
| PeopleSemSegNet AMR | Semantic segmentation of persons in an image. | nvidia/tao/peoplesemsegnet_amr |
| PCB Defect Classification | Model to classify defects in soldered components on a Printed Circuit Board. | nvidia/tao/pcb_classification |
| Gaze Estimation | Detect a persons eye gaze point of regard and gaze vector. | nvidia/tao/gazenet |
| Model name | Description | **NGC Instance |
|---|---|---|
| TAO Pretrained EfficientDet-TF2 | Pretrained efficientnet backbones for TAO Toolkit's efficientdet-tf2 | nvidia/tao/pretrained_efficientdet_tf2 |
| Pre-trained Deformable DETR ImageNet weights | Pre-trained deformable_detr weights trained on ImageNet to facilitate transfer learning using TAO Toolkit. | nvidia/tao/pretrained_deformable_detr_imagenet |
| Pre-trained SegFormer ImageNet weights | Pre-trained SegFormer weights trained on ImageNet to facilitate transfer learning using TAO Toolkit. | nvidia/tao/pretrained_segformer_imagenet |
| Pre-trained SegFormer NvImageNet weights | Pre-trained SegFormer weights trained on NvImageNet to facilitate transfer learning using TAO Toolkit. | nvidia/tao/pretrained_segformer_nvimagenet |
| Pre-trained DINO ImageNet weights | Pre-trained DINO weights trained on ImageNet to facilitate transfer learning using TAO Toolkit. | nvidia/tao/pretrained_dino_imagenet |
| Pre-trained FAN based NVImageNet Classification weights | Pre-trained FAN weights trained on NVImageNet to facilitate transfer learning using TAO Toolkit. | nvidia/tao/pretrained_fan_classification_nvimagenet |
| Pre-trained ImageNet Back bones weights | Pre-trained weights trained on ImageNet to facilitate transfer learning using TAO Toolkit. | nvidia/tao/pretrained_imagenet_backbones |
| Pre-trained DINO NvImageNet weights | Pre-trained DINO weights trained on NvImageNet to facilitate transfer learning using TAO Toolkit. | nvidia/tao/pretrained_dino_nvimagenet |
| TAO Pretrained Classification-TF2 | Pretrained backbones for TAO Toolkit TF2 image classification | nvidia/tao/pretrained_classification_tf2 |
| Pre-trained Segformer - CityScapes | Pre-trained segformer models trained on CityScapes. | nvidia/tao/pretrained_segformer_cityscapes |
| Pre-trained GCViT NVImageNet Classification weights | Pre-trained GcViT weights trained on NVImageNet to facilitate transfer learning using TAO Toolkit. | nvidia/tao/pretrained_gcvit_classification_nvimagenet |
| TAO Pretrained EfficientDet | Pretrained weights to facilitate transfer learning using TAO Toolkit. | nvidia/tao/pretrained_efficientdet |
| TAO Pretrained Semantic Segmentation | Pretrained weights to facilitate transfer learning using Transfer Learning Toolkit. | nvidia/tao/pretrained_semantic_segmentation |
| Pre-trained FAN based ImageNet Classification weights | Pre-trained FAN weights trained on ImageNet to facilitate transfer learning using TAO Toolkit. | nvidia/tao/pretrained_fan_classification_imagenet |
| Pre-trained Deformable DETR NvImageNet weights | Pre-trained deformable_detr weights trained on NvImageNet to facilitate transfer learning using TAO Toolkit. | nvidia/tao/pretrained_deformable_detr_nvimagenet |
| Pre-trained GCViT ImageNet Classification weights | Pre-trained GCViT weights trained on ImageNet to facilitate transfer learning using TAO Toolkit. | nvidia/tao/pretrained_gcvit_classification_imagenet |
| Pre-trained NVImageNet Back bones weights | Pre-trained weights trained on NVImageNet to facilitate transfer learning using TAO Toolkit. | nvidia/tao/pretrained_nvimagenet_backbones |
| Pre-trained EfficientDet Model trained on COCO | Pre-trained EfficientDet models trained on COCO to facilitate transfer learning using TAO Toolkit. | nvidia/tao/pretrained_efficientdet_tf2_coco |
| Pretrained Mask Auto Label | Pretrained model to facilitate transfer learning for MAL on TAO Toolkit | nvidia/tao/pretrained_mask_auto_label |
| TAO Pretrained Object Detection | Pretrained weights to facilitate transfer learning using TAO Toolkit. | nvidia/tao/pretrained_object_detection |
| TAO Pretrained DetectNet V2 | Pretrained weights to facilitate transfer learning using TAO Toolkit. | nvidia/tao/pretrained_detectnet_v2 |
| Pre-trained EfficientDet NvImageNet backbones | Pre-trained EfficientNet bacbone weights trained on NvImageNet to facilitate transfer learning using TAO Toolkit. | nvidia/tao/pretrained_efficientdet_tf2_nvimagenet |
| TAO Pretrained Instance Segmentation | Pretrained weights to facilitate transfer learning using TAO Toolkit. | nvidia/tao/pretrained_instance_segmentation |
| TAO Pretrained Classification | Pretrained weights to facilitate transfer learning using TAO Toolkit. | nvidia/tao/pretrained_classification |
Requirements
Recommended Hardware requirements
The following system configuration is recommended to achieve reasonable training performance with TAO Toolkit and supported models provided:
- 16 GB system RAM
- 16 GB of GPU RAM
- 8 core CPU
- 1 NVIDIA GPU
- 100 GB of SSD space
TAO Toolkit is supported on discrete GPUs, such as H100, A100, A40, A30, A2, A16, A100x, A30x, V100, T4, Titan-RTX and Quadro-RTX.
Note: TAO Toolkit is not supported on GPU's before the Pascal generation
Software requirements
| Software | Version | Comment |
|---|---|---|
| Ubuntu LTS | 20.04 | |
| python | >=3.6.9<3.7 | Not needed if you are using TAO API (See #3 below) |
| docker-ce | >19.03.5 | Not needed if you are using TAO API (See #3 below) |
| docker-API | 1.40 | Not needed if you are using TAO API (See #3 below) |
nvidia-container-toolkit |
>1.3.0-1 | Not needed if you are using TAO API (See #3 below) |
| nvidia-container-runtime | 3.4.0-1 | Not needed if you are using TAO API (See #3 below) |
| nvidia-docker2 | 2.5.0-1 | Not needed if you are using TAO API (See #3 below) |
| nvidia-driver | >525.85 | Not needed if you are using TAO API (See #3 below) |
| python-pip | >21.06 | Not needed if you are using TAO API (See #3 below) |
Package Content
Download the TAO package which contains startup scripts, Jupyter notebooks and config files.
TAO is supported on Google Colab; if you want to try on Colab, you can skip this step and directly scroll down to #4 in the How to run TAO section.
ngc registry resource download-version nvidia/tao/tao-getting-started:5.3.0 --dest ./
cd ./getting_started_v5.3.0
File Hierarchy
setup
|--> quickstart_launcher.sh
|--> quickstart_api_bare_metal
|--> quickstart_api_aws_eks
|--> quickstart_api_azure_aks
|--> quickstart_api_gcp_gke
notebooks
|--> tao_api_starter_kit
|--> api
|--> automl
|--> end2end
|--> dataset_prepare
|--> client
|--> automl
|--> end2end
|--> dataset_prepare
|--> tao_launcher_starter_kit
|--> dino
|--> deformable_detr
|--> classification_pyt
|--> ocdnet
|--> ...
|--> tao_data_services
|--> data
|--> ...
How to run TAO?
TAO is available as a docker container or as a collection of Python wheels.
There are 4 ways to run TAO depending on user preference and their setup. See the full list below.
1. Launcher CLI
The TAO Launcher is a lightweight Python based CLI application to run TAO. The launcher basically acts as a front-end for the multiple TAO Toolkit containers built on both PyTorch and Tensorflow. The multiple containers essentially get launched automatically based on the type of model you plan to use for your computer vision or conversational AI use-cases.
TAO Launcher
To get started, use the setup/quickstart_launcher.sh to validate your setup and install TAO launcher. Jupyter notebooks to train using the Launcher is provided under notebooks/launcher_starter_kit.
Detail instructions on installing pre-requisite and setup is provided in TAO documentation - Launcher
2. Directly from Container
Users have option to also run TAO directly using the docker container. To use container directly, user needs to know which container to pull. There are multiple containers under TAO, and depending on the model that you want to train you will need to pull the appropriate container. This is not required when using the Launcher CLI.
export DOCKER_REGISTRY="nvcr.io"
export DOCKER_NAME="nvidia/tao/tao-toolkit"
export DOCKER_TAG="***" ## for TensorFlow docker
export DOCKER_TAG="***" ## for PyTorch docker
export DOCKER_CONTAINER=$DOCKER_REGISTRY/$DOCKER_NAME:$DOCKER_TAG
docker run -it --rm --gpus all -v /path/in/host:/path/in/docker $DOCKER_CONTAINER \
detectnet_v2 train -e /path/to/experiment/spec.txt -r /path/to/results/dir -k $KEY --gpus 4
More information about running directly from docker is provided in TAO documentation - Container
3. TAO APIs
TAO Toolkit API is a Kubernetes service that enables building end-to-end AI models using REST APIs. The API service can be installed on a Kubernetes cluster (local / AWS EKS) using a Helm chart along with minimal dependencies. TAO toolkit jobs can be run using GPUs available on the cluster and can scale to a multi-node setting. Users can use a TAO client CLI to interact with TAO services remotely or can integrate it in their own apps and services directly using REST APIs.
TAO API
To get started, use the provided one-click deploy script to deploy either on bare-metal setup or on managed Kubernetes service like Amazon EKS. Jupyter notebooks to train using the APIs directly or using the client app is provided under notebooks/api_starter_kit
setup/quickstart_api_bare_metal
setup/quickstart_api_aws_eks
More information about setting up the API services and the API is provided in TAO documentation - API
4. Python Wheel
Users can also run TAO directly on bare-metal without docker or K8s. Users can deploy TAO notebooks directly on Google Colab without having to configure infrastructure. The full instructions are provided in the Colab notebook below.
| CV Task | Model Arch | One-click Deploy |
|---|---|---|
| Classification | ResNet18 | Train on Colab |
| Multi-task Classification | ResNet18 | Train on Colab |
| Object Detection | DSSD | Train on Colab |
| Object Detection | EfficientDet | Train on Colab |
| Object Detection | RetinaNet | Train on Colab |
| Object Detection | SSD | Train on Colab |
| Object Detection | YOLOv3 | Train on Colab |
| Object Detection | YOLOv4 | Train on Colab |
| Object Detection | YoloV4 Tiny | Train on Colab |
| Action Recognition | ActionRecognition | Train on Colab |
| OCR | LPRNet | Train on Colab |
| Pose Action Classification | PoseClassificationNet | Train on Colab |
| Emotion Recognition | EmotionNet | Train on Colab |
| Gesture Recognition | GestureNet | Train on Colab |
| Heart Rate Estimation | HeartRateNet | Train on Colab |
After starting TAO service locally or remotely, start Jupyter notebook
jupyter notebook --ip 0.0.0.0 --port 8888 --allow-root
Open an internet browser on localhost and navigate to the following URL: http://0.0.0.0:8888
Open the notebook that you are interested in training and start training.
Note: All the instructions to train, prune, optimize and download pretrained models are provided in the notebook.
Jupyter notebooks
All Notebooks and required spec files are provided in this package. The table below maps which notebook to use for fine-tuning either a purpose-build models like PeopleNet or an open model architecture like YOLO.
| Purpose-built Model | Launcher CLI notebook |
|---|---|
| PCB Classification | notebooks/tao_launcher_starter_kit/classification_pyt/classification.ipynb |
| Retail Object Recognition | notebooks/tao_launcher_starter_kit/metric_learning_recognition/metric_learning_recognition.ipynb |
| Optical Inspection | notebook/tao_launcher_starter_kit/optical_inspection/OpticalInspection.ipynb |
| Mask Auto Label | notebooks/tao_launcher_starter_kit/mal/mal.ipynb |
| OCRNet | notebooks/tao_launcher_starter_kit/ocrnet/ocrnet.ipynb |
| OCDNet | notebooks/tao_launcher_starter_kit/ocdnet/ocdnet.ipynb |
| PeopleSemSegFormer | notebooks/tao_launcher_starter_kit/segformer/segformer.ipynb |
| PeopleNet | notebooks/tao_launcher_starter_kit/detectnet_v2/detectnet_v2.ipynb |
| TrafficCamNet | notebooks/tao_launcher_starter_kit/detectnet_v2/detectnet_v2.ipynb |
| DashCamNet | notebooks/tao_launcher_starter_kit/detectnet_v2/detectnet_v2.ipynb |
| FaceDetectIR | notebooks/tao_launcher_starter_kit/detectnet_v2/detectnet_v2.ipynb |
| VehicleMakeNet | notebooks/tao_launcher_starter_kit/classification/classification.ipynb |
| VehicleTypeNet | notebooks/tao_launcher_starter_kit/classification/classification.ipynb |
| PeopleSegNet | notebooks/tao_launcher_starter_kit/mask_rcnn/mask_rcnn.ipynb |
| PeopleSemSegNet | notebooks/tao_launcher_starter_kit/unet/unet_isbi.ipynb |
| Bodypose Estimation | notebooks/tao_launcher_starter_kit/bpnet/bpnet.ipynb |
| License Plate Detection | notebooks/tao_launcher_starter_kit/detectnet_v2/detectnet_v2.ipynb |
| License Plate Recognition | notebooks/tao_launcher_starter_kit/lprnet/lprnet.ipynb |
| Facial Landmark | notebooks/tao_launcher_starter_kit/fpenet/fpenet.ipynb |
| FaceDetect | notebooks/tao_launcher_starter_kit/facenet/facenet.ipynb |
| ActionRecognitionNet | notebooks/tao_launcher_starter_kit/action_recognition_net/actionrecognitionnet.ipynb |
| PoseClassificationNet | notebooks/tao_launcher_starter_kit/pose_classification_net/pose_classificationnet.ipynb |
| ReIdentificationNet | notebooks/tao_launcher_starter_kit/re_identification_net/reidentificationnet.ipynb |
| Visual ChangeNet Classification | notebooks/tao_launcher_starter_kit/visual_changenet/visual_changenet_classification.ipynb |
| Visual ChangeNet Segmentation - Research Only (LevirCD) | notebooks/tao_launcher_starter_kit/visual_changenet/visual_changenet_segmentation.ipynb |
| Visual ChangeNet Segmentation - Commercial (LandSATSCD) | notebooks/tao_launcher_starter_kit/visual_changenet/visual_changenet_segmentation.ipynb |
| Open model architecture | Jupyter notebook |
|---|---|
| Deformable DETR | notebooks/tao_launcher_starter_kit/deformable_detr/deformable_detr.ipynb |
| DINO | notebooks/tao_launcher_starter_kit/dino/dino.ipynb |
| Image Classification | notebooks/tao_launcher_starter_kit/classification_pyt/classification_pyt.ipynb |
| Image Classification | notebooks/tao_launcher_starter_kit/classification_tf2/classification.ipynb |
| Image Classification | notebooks/tao_launcher_starter_kit/classification_tf1/classification.ipynb |
| Optical Inspection | notebooks/tao_launcher_starter_kit/optical_inspection/optical_inspection.ipynb |
| Metric Learning Recognition | notebooks/tao_launcher_starter_kit/metric_learning_recognition/metric_learning_recognition.ipynb |
| Segformer | notebooks/tao_launcher_starter_kit/segformer/segformer.ipynb |
| DetectNet_v2 | notebooks/tao_launcher_starter_kit/detectnet_v2/detectnet_v2.ipynb |
| FasterRCNN | notebooks/tao_launcher_starter_kit/faster_rcnn/faster_rcnn.ipynb |
| YOLOV3 | notebooks/tao_launcher_starter_kit/yolo_v3/yolo_v3.ipynb |
| YOLOV4 | notebooks/tao_launcher_starter_kit/yolo_v4/yolo_v4.ipynb |
| YOLOv4-Tiny | notebooks/tao_launcher_starter_kit/yolo_v4_tiny/yolo_v4_tiny.ipynb |
| SSD | notebooks/tao_launcher_starter_kit/ssd/ssd.ipynb |
| DSSD | notebooks/tao_launcher_starter_kit/dssd/dssd.ipynb |
| RetinaNet | notebooks/tao_launcher_starter_kit/retinanet/retinanet.ipynb |
| MaskRCNN | notebooks/tao_launcher_starter_kit/mask_rcnn/mask_rcnn.ipynb |
| UNET | notebooks/tao_launcher_starter_kit/unet/unet_isbi.ipynb |
| EfficientDet | notebooks/tao_launcher_starter_kit/efficientdet/efficientdet.ipynb |
| Mask Auto Label | notebooks/tao_launcher_starter_kit/mal/mal.ipynb |
| Visual ChangeNet Classification | notebooks/tao_launcher_starter_kit/visual_changenet/visual_changenet_classification.ipynb |
| Visual ChangeNet Segmentation | notebooks/tao_launcher_starter_kit/visual_changenet/visual_changenet_segmentation.ipynb |
Important Links
Blogs
Train like a 'pro' with AutoML in TAO
Deploy TAO on Azure ML
Synthetic Data and TAO
Action Recognition Blog
Real-time License Plate Detection
2 Pose Estimation: Part 1
Part 2
Building ConvAI with TAO Toolkit
License
TAO Toolkit getting Started License for TAO containers is included in the banner of the container. License for the pre-trained models are available with the model cards on NGC. By pulling and using the Train Adapt Optimize (TAO) Toolkit container to download models, you accept the terms and conditions of these licenses.