5 steps to implement deep learning on Jetson Nano

Welcome! If you are planning to dive into edge computing and deploy deep learning models on embedded hardware, you have come to the right place.

The NVIDIA Jetson Nano Developer Kit is a highly capable embedded computing board designed to run machine learning algorithms in parallel. It is an excellent, low-power platform (running on as little as 5 volts) for executing complex tasks like image classification, object detection, segmentation, and speech processing.

In this guide, I will walk you through the essential steps to prepare your Jetson Nano environment, from installing the correct operating system to setting up TensorFlow with GPU acceleration.

While this tutorial references the official NVIDIA documentation, I have structured it to provide a streamlined, battle-tested path. Following this guide will help you avoid the common library conflicts and compatibility issues that beginners often face when setting up their boards for the first time.

1. Flashing the Operating System

To ensure optimal performance and adequate storage for your datasets and models, I strongly recommend using a high-speed microSD card with at least a 32 GB capacity.

You can follow the official getting started guide to flash the OS. However, it is crucial that you do not install a generic Ubuntu image. Instead, you must install the NVIDIA-provided image that comes pre-packaged with the CUDA-X accelerated libraries, known as JetPack.

Although newer versions exist, for maximum stability and compatibility with specific TensorFlow distributions, I highly recommend installing JetPack 4.4.1.
Download the image here: JetPack SDK 4.4.1 Archive

Hardware Note: Unlike the Raspberry Pi, the standard Jetson Nano Developer Kit does not feature an onboard Wi-Fi module. Ensure you have an Ethernet cable ready, or purchase a compatible M.2 Wi-Fi card or USB Wi-Fi dongle before beginning the setup process.

Once you boot up successfully, you should see the Ubuntu desktop environment. A directory named L4T-README will be present on your desktop, confirming a successful JetPack installation.

2. Installing System Prerequisites and Dependencies

Before installing TensorFlow, we need to ensure the system is up-to-date and equipped with the necessary compiler tools and mathematical libraries. Open your terminal and execute the following commands:

Update system repositories:

sudo apt-get update

Install Python package manager (PIP):

sudo apt-get install python3-pip

Install core system libraries (HDF5, BLAS, LAPACK):

sudo apt-get install libhdf5-serial-dev hdf5-tools libhdf5-dev 
sudo apt-get install zlib1g-dev zip libjpeg8-dev liblapack-dev libblas-dev gfortran

Install Python dependencies required by TensorFlow:

sudo pip3 install -U pip testresources setuptools 
sudo pip3 install -U numpy==1.16.1 future==0.17.1 mock==3.0.5 h5py==2.9.0
sudo pip3 install -U pip keras_preprocessing==1.0.5 keras_applications==1.0.8
sudo pip3 install -U gast==0.2.2 futures protobuf pybind11

3. Installing TensorFlow with GPU Support

Because the Jetson Nano utilizes an ARM architecture alongside an NVIDIA Maxwell GPU, standard TensorFlow binaries from PyPI will not work. We must install the specific pre-compiled wheels hosted by NVIDIA.

To install TensorFlow 1.15.0, run:

sudo pip3 install --pre --extra-index-url https://developer.download.nvidia.com/compute/redist/jp/v441 'tensorflow<2'

Alternatively, if your project requires TensorFlow 2.x, you can specify the version like this (e.g., TensorFlow 2.4.0):

sudo pip3 install --extra-index-url https://developer.download.nvidia.com/compute/redist/jp/v441 'tensorflow==2.4.0+nv21.5'

4. Verifying the Setup

To confirm that TensorFlow is installed correctly and can successfully initialize the GPU, run this one-liner in your terminal:

python3 -c 'import tensorflow as tf; print(tf.__version__)'

If the installation was successful, the terminal will print your currently installed TensorFlow version without any critical library errors.

5. Customizing Your Workspace (Neofetch)

As a developer, customizing your terminal can make your workflow more enjoyable. Neofetch is a command-line utility that displays your operating system logo alongside vital system information.

Install it via apt:

sudo apt-get install neofetch
neofetch

To automatically display this dashboard every time you open a new terminal session, append the command to your bash profile:

sudo apt-get install nano
nano ~/.bashrc

Scroll to the very bottom of the file, type neofetch, save (Ctrl+O), and exit (Ctrl+X). Restart your terminal to see the effect in action.

6. Monitoring Performance with Jtop

When training or running inference on machine learning models, monitoring resource utilization is essential. While Linux users traditionally rely on htop, it cannot monitor NVIDIA's Tegra GPU metrics.

Instead, we use jtop (part of the jetson-stats package). It provides a real-time, comprehensive view of your CPU, GPU, memory usage, and thermal sensors.

Install jtop:

sudo -H pip install -U jetson-stats

Launch the monitor:

jtop

7. Running a Test Model

To validate that our entire pipeline—hardware, OS, and TensorFlow—is functioning in harmony, let's train a sample model using the classic MNIST dataset.

First, ensure Git is installed to clone the repository:

sudo apt-get install git

Clone the test script:

git clone https://github.com/redsriracha/tensorflow-gpu-test

Navigate into the directory and execute the script:

cd tensorflow-gpu-test
python3 run.py

If you see the epochs running and the loss decreasing, congratulations! Your NVIDIA Jetson Nano is successfully configured for deep learning.

A special thanks to Ngoc Truc for the support in creating this tutorial. If you have any questions or run into issues during the setup, feel free to drop a comment below!