Introduction: Transfer Learning With NVIDIA JetBot - Fun With Traffic Cones

About: I like AI and machine learning applications, especially in robotics

Teach your robot to find a path in a maze of traffic cones using the camera and state-of-the-art deep learning model.

Supplies

Step 1: Motivation

Every time I drive in the contraction area I am thinking how challenging it would be for a self-driving car to navigate through the traffic cones. It turns out it is not so difficult with new NVIDIA's JetBot - with only a couple hundred images, you can train a state-of-the-art deep learning model to teach your robot how to find a path in a maze of toy traffic cones using the onboard camera only and no other sensors.

Step 2: NVIDIA JetBot & Project Overview

JetBot is an open-source robot based on NVIDIA Jetson Nano kit. You can find detailed instructions on how to build and set it up here.

This project is a modified Collision avoidance example from NVIDIA JetBot Wiki. It consists of three major steps, each described in a separate Jupyter notebook:

  • Collect data on JetBot - notebook data_collection_cones.ipynb
  • Train model on other GPU machine - notebook train_model_cones.ipynb
  • Run live demo on JetBot - notebook live_demo_cones.ipynb

You can find these three Jupyter notebooks here

Step 3: Build JetBot and Upload Jupyter Notebooks

  1. Build and setup JetBot as explained here
  2. Connect to your robot by navigating to http://:8888Sign in with the default password jetbot
  3. Shutdown all other running notebooks by selecting Kernel -> Shutdown All Kernels...
  4. Navigate to ~/Notebooks/
  5. Create new subfolder ~/Notebooks/traffic_cones_driving/
  6. Upload data_collection_cones.ipynb and live_demo_cones.ipynb to ~/Notebooks/traffic_cones_driving/

IMPORTANT: The Jupyter notebooks data_collection_cones.ipynb and live_demo_cones.ipynb referenced in this instructions should be run on JetBot while train_model_cones.ipynb - on a computer with GPU.

Therefore we have to upload data_collection_cones.ipynb and live_demo_cones.ipynb to the JetBot and place them in ~/Notebooks/traffic_cones_driving/

Step 4: Collecting Training Data on JetBot

We will collect an image classification dataset that will be used to help JetBot operate in a maze of the traffic cones. JetBot will learn how to estimate probabilities of four scenarios (classes):

  • Free - when it is safe to move forward
  • Blocked - when there is an obstacle in front of the robot
  • Left - when the robot should spin to the left
  • Right - when the robot should spin to the right

To collect the training data on JetBot we will use Jupyter notebook data_collection_cones.ipynb which contains details instructions on how to do it. To run this notebook on JetBot, follow the next steps:

  1. Connect to your robot by navigating to http://:jetbot-ip-address::8888
  2. Sign in with the default password jetbot
  3. Shutdown all other running notebooks by selecting Kernel -> Shutdown All Kernels...
  4. Navigate to ~/Notebooks/traffic_cones_driving/
  5. Open and follow the data_collection_cones.ipynb notebook

Step 5: Train Neural Network on GPU Machine

Next, we will use the collected data to re-train the deep learning model AlexNet on GPU machine (host) by running train_model_cones.ipynb.

Note that train_model_cones.ipynb is the only Jupyter notebook in this tutorial which is run NOT on the JetBot.

  1. Connect to a GPU machine with PyTorch installed and a Jupyter Lab server running
  2. Upload train_model_cones.ipynb notebook and to this machine
  3. Upload dataset_cones.zip file that you created in the data_collection_cones.ipynb notebook and extract this dataset. (After this step you should see a folder named dataset_cones appear in the file browser.)
  4. Open and follow the train_model_cones.ipynb notebook. At the end of this step, you will create a model - the file best_model_cones.pth which then has to be uploaded to the JetBot to run the live demo.

Step 6: Run Live Demo on JetBot

This final step is to upload the model best_model_cones.pth to the JetBot and run it.

  1. Power your robot from the USB battery pack
  2. Connect back to your robot by navigating to http://:jetbot-ip-address::8888
  3. Sign in with the default password jetbot
  4. Shutdown all other running notebooks by selecting Kernel -> Shutdown All Kernels...
  5. Navigate to ~/Notebooks/traffic_cones_driving
  6. Open and follow the live_demo_cones.ipynb notebook

Start cautious and give JetBot enough space to move around. Try different cone configuration and see how well the robot performs in different environments, lighting, etc. While the notebook live_demo_cones.ipynb explains all steps in detail, the following chart shows the logic of robot moves given the probabilities predicted by the models.

The notebook also explains how to store the history of robot moves with free/left/right/blocked probabilities predicted by the model and how to make two FPV (First Person View) videos (at 1 fps and 15 fps rates) with superimposed telemetry and JetBot actions data. Those are useful for debugging, PID controller tuning and model improvement.


Have fun and let me know if you have questions! :-)

Code is available on Github