Skip to navigation

TAO Toolkit Integration with DeepStream

View as Markdown

NVIDIA TAO toolkit is a simple, easy-to-use training toolkit that requires minimal coding to create vision AI models using the user’s own data. Using TAO toolkit, users can transfer learn from NVIDIA pre-trained models to create their own model. Users can add new classes to an existing pre-trained model, or they can re-train the model to adapt to their use case. Users can use model pruning capability to reduce the overall size of the model.

Pre-trained models

There are 2 types of pre-trained models that users can start with - purpose-built pre-trained models and meta-architecture vision models. Purpose-built pre-trained models are highly accurate models that are trained on millions of objects for a specific task. The pre-trained weights for meta-architecture vision models merely act as a starting point to build more complex models. These pre-trained weights are trained on Open image dataset and they provide a much better starting point for training versus starting from scratch or starting from random weights. With the latter choice, users can choose from 100+ permutations of model architecture and backbone. See the illustration below.

TAO toolkit Pretrained Models

TAO toolkit pretrained models - use cases

Model NameDescriptionNGC InstanceTAO Finetuning
TrafficCamNet4 class object detection network to detect cars in an image.nvidia/tao/trafficcamnetYes
PeopleNet3 class object detection network to detect people in an image.nvidia/tao/peoplenetYes
VehicleMakeNetResnet18 model to classify a car crop into 1 out 20 car brands.nvidia/tao/vehiclemakenetYes
VehicleTypeNetResnet18 model to classify a car crop into 1 out 6 car types.nvidia/tao/vehicletypenetYes
License Plate DetectionObject Detection network to detect license plates in an image of a car.nvidia/tao/lpdnetYes
License Plate RecognitionModel to recognize characters from the image crop of a License Plate.nvidia/tao/lprnetYes
PeopleSemSegNetSemantic segmentation of persons in an image.nvidia/tao/peoplesemsegnetYes
CitySemSegFormerSemantic segmentation of persons in an image.nvidia/tao/citysemsegformerYes
Optical Character DetectionNetwork to detect characters in an image.nvidia/tao/ocdnetYes
Optical Character RecognitionModel to recognise characters from a preceding OCDNet model.nvidia/tao/ocrnetYes
Retail Object DetectionDINO (DETR with Improved DeNoising Anchor Boxes) based object detection network to detect retail objects on a checkout counter.nvidia/tao/retail_object_detectionYes
Retail Object RecognitionEmbedding generator model to recognize objects on a checkout counter.nvidia/tao/retail_object_recognitionYes
ReIdentificationNetRe-Identification network to generate embeddings for identifying persons in different scenes.nvidia/tao/reidentificationnetYes
Pose ClassificationPose classification network to classify poses of people from their skeletons.nvidia/tao/poseclassificationnetYes
PeopleNet Transformer3 class object detection network to detect people in an image.nvidia/tao/peoplenet_transformerYes
BodyPose3DNet3D human pose estimation network to predict 34 keypoints in 3D of a person in an image.nvidia/tao/bodypose3dnetYes
PointPillarNetObjects detection from a LIDAR point cloud file.nvidia/tao/pointpillarnetYes
Mask2FormerInstance segmentation network to predict pixel-accurate masks for detected objects in an image.nvidia/tao/mask2formerYes

Most models trained with TAO toolkit are natively integrated for inference with DeepStream. If the model is integrated, it is supported by the reference deepstream-app. If the model is not natively integrated in the SDK, you can find a reference application on the GitHub repo. See the table below for information on the models supported. For models integrated into deepstream-app, we have provided sample config files for each of the networks. The sample config files are available in the deepstream_tao_apps_configs folder. The table below also lists config files for each model.

Refer README in package /opt/nvidia/deepstream/deepstream/samples/configs/tao_pretrained_models/README.md to obtain TAO toolkit config files and models mentioned in following table.

The TAO toolkit pre-trained models table shows the deployment information of purpose-built pre-trained models.

TAO toolkit pre-trained models in DeepStream

Pre-Trained modelDeepStream reference appConfig filesDLA supported
TrafficCamNetdeepstream-appdeepstream_app_source1_trafficcamnet.txt, config_infer_primary_trafficcamnet.txt, labels_trafficnet.txtYes
PeopleNetdeepstream-appdeepstream_app_source1_peoplenet.txt, config_infer_primary_peoplenet.txt, labels_peoplenet.txtYes
PointPillarNetdeepstream-lidar-inference-appconfig_lidar_triton_infer.yamlNo
VehicleMakeNetdeepstream-appdeepstream_app_source1_dashcamnet_vehiclemakenet_vehicletypenet.txt, config_infer_secondary_vehiclemakenet.txt, labels_vehiclemakenet.txtYes
VehicleTypeNetdeepstream-appdeepstream_app_source1_dashcamnet_vehiclemakenet_vehicletypenet.txt, config_infer_secondary_vehicletypenet.txt, labels_vehicletypenet.txtYes
CitySemSegFormerdeepstream_tao_appsseg_app_unet.yml, pgie_citysemsegformer_tao_config.txtYes
License Plate Detectiondeepstream_tao_appsNo
License Plate Recognitiondeepstream_tao_appsNo
PoseClassificationNetdeepstream_tao_appsNo
BodyPose3DNetdeepstream_tao_appsNo
Retail Object Recognitiondeepstream_tao_appsNo
Retail Object Detectiondeepstream_tao_appsNo
PeopleNet Transformer Modeldeepstream_tao_appsNo
ReIdentificationNetdeepstream_tao_appsNo
Optical Character Detectiondeepstream_tao_appsNo
Optical Character Recognitiondeepstream_tao_appsNo
PeopleNet Transformerdeepstream_tao_appsdet_app_frcnn.yml, pgie_peoplenet_transformer_tao_config.ymlNo
Mask2Formerdeepstream_tao_appsins_seg_app.yml, pgie_mask2former_tao_config.ymlNo

The TAO toolkit models can be deployed with DeepStream sample applications.

Learn more about running the TAO models here: deepstream tao apps

For more information about TAO and how to deploy TAO models, refer to Integrating TAO Models into DeepStream chapter of TAO toolkit user guide.

For more information about deployment of architecture specific models with DeepStream, refer to deepstream tao apps Github repo.