Retrain SSD MobileNet Object Detection Model with Pytorch on Colab
Transfer Learning with Pytorch
Transfer learning is a technique for re-training a DNN model on a new dataset, which takes less time than training a network from scratch. With transfer learning, the weights of a pre-trained model are fine-tuned to classify a customized dataset. In these examples, we'll be using the ResNet-18 and SSD-Mobilenet networks, although you can experiment with other networks too.
PyTorch is the machine learning framework that we'll be using, and example datasets along with training scripts are provided to use below, in addition to a camera-based tool for collecting and labeling your own training datasets.
Installing Pytorch on NVIDIA Jetson
If you are Running the Docker Container or optionally chose to install PyTorch back when you Built the Project, it should already be installed on your Jetson to use. Otherwise, if you aren't using the container and want to proceed with transfer learning, you can install it now:$ cd jetson-inference/build
$ ./install-pytorch.sh
Verifying Pytorch
You can test that PyTorch was installed correctly and detects your GPU by executing these commands from an interactive Python shell - run python or python3 from a terminal:
>>> import torchvision >>> torchvision.__version__
Re-training SSD-Mobilenet
Next, we'll train our own SSD-Mobilenet object detection model using PyTorch and the Open Images dataset. SSD-Mobilenet is a popular network architecture for realtime object detection on mobile and embedded devices that combines the SSD-300 Single-Shot MultiBox Detector with a Mobilenet backbone.
In the example below, we'll train a custom detection model that locates 8 different varieties of fruit, although you are welcome to pick from any of the 600 classes in the Open Images dataset to train your model on. You can visually browse the dataset here.
Downloading the Dataset
The Open Images dataset contains over 600 object classes that you can pick and choose from. There is a script provided called open_images_downloader.py which will automatically download the desired object classes for you.
The classes that we'll be using are "Apple,Orange,Banana,Strawberry,Grape,Pear,Pineapple,Watermelon", for example for a fruit-picking robot - although you are welcome to substitute your own choices from the class list. The fruit classes have ~6500 images, which is a happy medium.
Retrain model Diagram
Retrain with Pytorch. ( colab )
Then convert pytorch model to ONNX model ( colab )
then Run ONNX model and convert to TensorRT model on NVIDIA Jetson
Retrain Object Detection Model on Colab
Run Jetson Inference with detectnet
./detectnet --model=networks/fruit/ssd-mobilenet.onnx \
--labels=networks/fruit/labels.txt --input-blob=input_0 \
--output-cvg=scores --output-bbox=boxes \
images/fruit_0.jpg fruit0.jpg
PART-2 Run Model on NVIDIA Jetson
Reference
Jetson Inference. https://github.com/dusty-nv/jetson-inference
Open Images Dataset https://storage.googleapis.com/openimages/web/index.html
