Custom Image Classification with Google Coral TPU USB Accelerator

In this blog post, I’ll cover making predictions on our custom image datasets using the Google Coral TPU USB Accelerator. First of all, let’s take a quick look at what the Tensor Processing Unit is and how it works.

TPUs can be defined as custom-developed application-specific integrated circuits for machine learning tasks. They involve MatriX multiplier Units (MXU) and Vector Processing Units (VPU). As you can understand, the Matrix Multiplier Unit handles matrix multiplications. This unit uses an architecture called a systolic array. Systolic array architecture is a chain of Arithmetic Logic Units (ALUs) that are connected to each other. Those ALUs only perform multiplications and additions.

Figure 1: Systolic Array Architecture

Vector Processing Units handle all other tasks such as activations, softmax, etc.

If you are interested in processor/computer architecture, you can find a Verilog implementation of a simple TPU here.

Google’s Edge TPU USB accelerator is a coprocessor that you can connect via USB to your computer. According to the official website, the accelerator is capable of performing 4 trillion operations (tera-operations) per second (TOPS), using 0.5 watts for each TOPS (2 TOPS per watt).

Figure 2: Google Coral Edge TPU

You need to install some libraries to use the Edge TPU. It supports Windows, macOS, and Linux distros. You can find instructions here to install those libraries on your computer.

Once you installed the libraries, you are ready to train a model, convert it to TensorFlow Lite format, and make predictions. You can learn how to train a custom image classification model and convert it into the .tflite format in my previous blog post.

I use a reduced and balanced version of the Dermnet dataset. There are 4 classes and each of those classes contains 1040 samples. For under-represented classes, I used data augmentation techniques, which you can find the code in my previous blog post. To run our model on the Edge TPU accelerator, a few extra steps are required:

1. The original .h5 model requires post-training quantization. That document can be found here. You can quantize your model by following this document.
2. The quantized model must then be compiled for use with the Edge TPU. Here is a document for it.

You can download the Edge TPU model from here.

The Edge TPU, on the other hand, does not support all operations. A list of supported operations can be found here.

I’ll proceed with this model to show how to make predictions using the Edge TPU USB Accelerator.

I simply connect the device via USB and run this script.

Make sure that the script, the model, labels.txt file, and an image you picked from the test set (in this case, it’s acute-paronychia-7.jpg) are in the same directory. Run the script by typing:

>python predict.py

You’ll see something like this:

Figure 3: Results

Making Predictions Using Models with Custom Layers

You can write your own custom Keras layers to use on your model, like below. This example is a simple model with just one layer. You can convert it to the .tflite file and make predictions with the Google Coral Edge TPU.

You’ll see something like this:

Figure 4: Results for custom layer model

Through this article, you have learned how to create a vision-based ML project with TensorFlow Lite using the Google Coral TPU USB Accelerator. Following that you learned how to use that model on device for machine learning on the edge.