Installing Jax on the Jetson Nano

By -
Yesterday I got Jax running on a Jetson Nano. There's been some online interest and I promised to describe the installation. I'll cover the process below but first I'll briefly explain

What's Jax, and why do I want it?


tldr: The cool kids have moved from Tensorflow to Jax.
 
If you're interested in AI and Artificial Neural Networks (ANNs) you'll have heard of Google's TensorFlow.

Tensorflow with Keras makes it easy to
  • build an ANNs from standard software components
  • train the network
  • test it and
  • deploy it into production
Even better: TensorFlow can take advantage of a GPU or TPU if they are available, which allows a dramatic speedup of the training process. That's a big deal because training a complex network from scratch might take weeks of computer time.

However, I find it hard to get TensorFlow and its competitors to use novel network architectures.  I am currently exploring types of  network that I can't implement easily in TensorFlow.

I could code them in Python and Numpy but I want to be able to run them fast, and that means using a GPU.

Hence Jax.

You can read about why DeepMind uses Jax here. 

Jax implements many of Numpy's primitives on a GPU and it can also convert functional, side-effect-free Python code.

It also supports automatic differentiation. I'm not going to be using back propagation in the work I'm doing, so that's not a big deal for me. If you are, you can see online examples that build and train networks from scratch in a few lines of Python code.

I'll give some numbers for the speed-ups you get later in the article. For now, I'll just say that Jax code runs faster than Numpy on the Nano, and faster still when compared with Numpy on the Raspberry Pi.

Jax on the Jetson Nano

Jetson Nano with SSD

There are several ways you can run Jax-based applications; it's not hard to install on a Linux system, and you can run it on Google's wonderful, free Colab service. My interest lies in using neural networks with mobile robots, so I want to run it on an Edge Computing device with GPU hardware built-in.

The Jetson Nano is ideal for that.

I recently came across an article on NVIDIA's developer website that described how to install Jax on the Nano. It included a warning that the installation was slow and required a large (10GB) swap file.

I decided to re-configure one of my Nanos from scratch and set it up to boot from a USB SSD. I did the whole process in three stages using links which I have listed below, along with the gotcha I hit and my work-around.

  1. I set up the Nano in headless mode using an SD card.
  2. I configured the Nano to boot from a USB SSD.
  3. I installed Jax from source.

Set up the Nano in headless mode using an SD card


NVIDIA have really simplified the setup process since I started exploring the Nano a couple of years ago.

The NVIDIA Getting Started guide now gives two ways to set up your Nano.

One requires an HDMI screen, a mouse and keyboard. 

The other needs an Ethernet connection, a suitable power supply and a jumper to configure the Nano to power itself from its barrel connector rather than from USB.

I want to run my Nano in fast, full-power mode and I have very limited desk space. I decided to go for a Headless Setup. I followed these three steps in the NVIDIA guide:

  1. Prepare for Setup*
  2. Write Image to the microSD Card
  3. Setup and First Boot**
* There are alternative power supplies to the Adafruit supply that NVIDIA mentioned. I purchased this one in the UK.

** Make sure you scroll down to the section Initial Setup Headless Mode

Configure the Nano to boot from USB SDD


I followed this excellent guide from JetsonHacks.

I found that the Nano did not boot from the SSD first time but it's booted reliably ever since.

Installing Jax from source

I followed the process given in this post on the NVIDIA developer website. I hit one snag, which was easy to solve: After I had set up a Python 3.9 virtual environment, I needed to run

sudo apt install python3.9-distutils

before I could run the stages that used pip.

First time through I failed to realise that the heading How to Add Swap Space on Ubuntu 18.04 is actually a link to a web page that tells you how to do that. I also set the swap file size to 16 GB rather than 10GB as recommended.

The installation guide says that the compilation process takes 12 hours. Mine took about 6 hours, possibly because I was running the Nano in Max Power mode.

Once the compilation process has finished, the remaining steps take a minute or less!

First Timing Tests

It looks as if Jax is six to seven times faster than numpy on the Nano when dot-multiplying a 4000x4000 element matrix with itself, and 25 times faster than a Pi 4 running numpy. That's worth the effort of installation :)



 




Comments

Post a Comment

Popular posts from this blog

Controlling a Raspberry Pi Pico remotely using PySerial

By -
Image
Update : I'm using thie code below in another project, and found that I had not correctly fixed the reported bug. The new version passes automated tests, and I am pretty sure it works OK. I have changed the name of the class to Talker  since it can both send and receive information. Apologies to all concerned for the bug! Introduction You can use a Raspberry Pi Pico as a powerful peripheral to a host - a Raspberry Pi, a Jetson Nano, a laptop or workstation. In this article you'll see how to interact with a Pico running MicroPython or CircuitPython by writing some Python code that runs on the host. The software is easy to use. It enables you to send a Python statement to the Pico and read the results. The statement can be any valid MicroPython code. Setting up the host and the Pico For this article I've used a Raspberry Pi as the host, but any computer running Windows, Linux or Mac OS will do so long as it has Python 3.5 or later installed. In particular, you can use this t...
Read more

Five steps to connect Jetson Nano and Arduino

By -
Image
Yesterday's post showed how to link a micro:bit to the Jetson Nano. One of the members of the (unofficial) NVIDIA Jetson Nano group on Facebook asked about connecting an Arduino to the Jetson. Here's a simple recipe for getting data from the Arduino to the Jetson Nano. It should work on all the Jetson models, not just the Nano, but I only have Nanos to hand. On request, I've added a recipe at the end of this post which sends data from the Jetson Nano to the Arduino; it turns the default LED on the Arduino on or off. The recipe for sending data from the Arduino to the Jetson has just 5 stages: Program the Arduino. (I used the ASCIITable example). Connect the Arduino to the Jetson using a USB connector Install pyserial on the Jetson Download a three-line Python script Run the script.   Programming the Arduino I used an Arduino Uno, and checked it on a verteran Duemilanove (above), but any Arduino should work. You'll need to do this step usi...
Read more

Raspberry Pi Pico project 2 - MCP3008

By -
Image
This post show you how to drive the MCP3008 8-channel ADC using a Raspberry Pi Pico. I'll start with a personally embarrassing and annoyingly relevant story. I'll also tell you about a minor 'gotcha' when using SPI on the Pico, and how you can avoid it. Finally, I've include the code and fritzing diagram I used to test the interface. First, the story. Back in 2012, when the Raspberry Pi was fresh and new, I was running a startup called Quick2Wire.  We made add-on boards for the Pi. Just before Christmas we sent out our first batch of boards to a small army of beta-testers. The beta boards didn't work. Somehow a single trace on the pcb had got deleted and one of the GPIO pins was isolated. The boards were still usable and it wasn't too hard to solder in a jumper wire to replace the missing trace, but it was very embarrassing. It wasn't just that we'd shipped boards with a defect. The really annoying thing was that I'd set up an  automated loop-bac...
Read more