This repo is about how to train a custom object detection model (using DetectNet_v2 architecture) on our own dataset using Nvidia Transfer Learning Toolkit (TLT-kit) and AWS GPU instance. Later, using the trained model to do inferencing on Jetson nano device with the help of Nvidia Deepstream SDK.
- The NVIDIA TLT-kit is used with NVIDIA pre-trained models to create custom Computer Vision (CV) and Conversational AI models with the user’s own data. Training AI models using TLT-kit does not require expertise in AI or deep learning.
- NVIDIA's DeepStream SDK delivers a complete streaming analytics toolkit for AI-based multi-sensor processing, video, audio and image understanding. DeepStream is for vision AI developers, software partners, startups and OEMs building IVA apps and services.
- Hands-on practice in AWS instance.
- AWS GPU instance (g4dn.xlarge) as a training machine.
- An NGC account to download Nvidia pre-trained models.
-
Connect to an AWS GPU instance.
-
Make nested directories
mkdir -p tlt_experiments/data tlt_experiments/detectnet_v2 -
Move the dataset to the
tlt_experiments/data/folder. Make sure the dataset is in KITTI format. Follow this to make your dataset in KITTI format. -
Set up an NGC account, generate API KEY and download Command Line Interface (CLI)
- Create an NGC account here and sign in.
- Generate an API KEY here and save the KEY somewhere.
- Download CLI by following the instructions here.
For more information follow this.
-
Set up a TLT container for training. Optionally, upgrade the python version if required from here. Also make sure the port 8888 is enabled in the EC2 instance.
-
Pull TLT-kit docker image using
docker pull nvcr.io/nvidia/tlt-streamanalytics:v2.0_py3 -
Run docker container using
docker run --runtime=nvidia -it -v /home/ubuntu/tlt-experiments:/workspace/tlt-experiments -p 8888:8888 nvcr.io/nvidia/tlt-streamanalytics:v2.0_py3 /bin/bash -
You will end up in the container.
-
-
To start the container later follow this.
-
To start the TLT container
docker start <container id> -
To attach the TLT container
docker attach <container id>
-
-
Install TLT-kit and check if it’s successfully installed.
-
Install TLT-kit by running the following.
pip3 install nvidia-pyindexpip3 install nvidia-tlt -
Check if it’s successfully installed.
tlt --help
-
-
(Optionally) Install and configure the jupyter notebook if it's not already by following this.
-
Open jupyter notebook
-
To open the jupyter notebook run the following command.
jupyter notebook --ip 0.0.0.0 --no-browser --port 8888 --allow-root -
Copy-paste the URL to the browser and replace the localhost IP in the URL with the public IP of the EC2 instance.
-
-
Goto
examples/detectnet_v2directory and run the notebook. Step-by-step follow detectnet_v2.ipynb notebook for training. Follow this for TLT kit documentation. -
While exporting the model for Deepstream on a Jetson nano device make sure it’s FP32 or FP16. Jetson nano device doesn’t support INT8.
-
Once you get the .etlt model, generate a labels.txt file containing the class names. Follow instructions from here for more details.
-
Move .etlt model and labels.txt file in deepstream-python-apps directory on Jetson nano device.
-
Use deepstream-test3 app and configure the dstest3_pgie_config.txt by following the instructions from here. Set API KEY in the dstest3_pgie_config.txt as
tlt-model-key=tlt_encode -
Run deepstream-test3 by following this and see the output.
I have trained an object detection model to detect the defect on the surface of metal objects. This type of models can be used in any manufacturing industries for quality check purpose of metal parts. Here is the inference output demo of our Defect Detection model.
- Follow me on Linkedin: mohdsaqibhbi
- Subscribe my Youtube channel: StarrAI