Run inference on the Edge TPU with Python

To simplify development with Coral, we made the Edge TPU compatible with the standard TensorFlow Lite API. So if you already have code using the TensorFlow Lite API, then you can run a model on the Edge TPU by changing just a couple lines of code.

And to make development even easier, we created a wrapper library—the PyCoral API—to handle a lot of boilerplate code that's required when running an inference with TensorFlow Lite. Plus, this library provides features specific for the Edge TPU, such as model pipelining and on-device transfer learning.

In the following sections, we'll show you how to run an inference using the PyCoral API or update your existing TensorFlow Lite code to work with the Edge TPU.

Overview

By default, the TensorFlow Lite interpreter executes your model on the CPU, but that fails if your model is compiled for the Edge TPU because an Edge TPU model includes a custom Edge TPU operator. So you need to specify a delegate for the Edge TPU. Then, whenever the interpreter encounters the Edge TPU custom operator, it sends that operation to the Edge TPU. As such, inferencing on the Edge TPU requires only the TensorFlow Lite API.

So if you already have code that runs a TensorFlow Lite model, you can update it to run your model on the Edge TPU by following the steps below to update your TensorFlow Lite code.

Otherwise, we recommend you try the PyCoral API, because it hides a lot of repetitive TensorFlow Lite code behind simple functions and offers some additional features for the Edge TPU.

Run an inference with the PyCoral API

The PyCoral API is a small set of convenience functions that initialize the TensorFlow Lite Interpreter with the Edge TPU delegate and perform other inferencing tasks such as parse a labels file, pre-process input tensors, and post-process output tensors for common models.

To get started, all you need to do is the following:

1. Install the PyCoral API.

   • On Linux, you can install it with a Debian package:

     sudo apt-get update
     
     
     sudo apt-get install python3-pycoral

   • On Mac or Windows, you can install the Python wheel.

2. In your Python code, import the pycoral module (or just some specific sub-modules).

3. Initialize the TensorFlow Lite Interpreter for the Edge TPU by calling make_interpreter().

4. Then use our "model adapter" functions to simplify your code—such as using classify.set_input() and classify.get_output() to process the input and output tensors—in combination with calls to the TensorFlow Lite Interpreter, as shown in the following example.

Again, the workflow is entirely based on the TensorFlow Lite interpreter API, so you should understand that API first. Then you can use our Coral API to remove some of the boilerplate code for initializing the interpreter and processing your input and output tensors.

For more detail about the APIs, see the PyCoral API reference.

Inferencing example

Just to show how simple your code can be, the following is a complete script that runs a classification model using the PyCoral API:

import os
import pathlib
from pycoral.utils import edgetpu
from pycoral.utils import dataset
from pycoral.adapters import common
from pycoral.adapters import classify
from PIL import Image

# Specify the TensorFlow model, labels, and image
script_dir = pathlib.Path(__file__).parent.absolute()
model_file = os.path.join(script_dir, 'mobilenet_v2_1.0_224_quant_edgetpu.tflite')
label_file = os.path.join(script_dir, 'imagenet_labels.txt')
image_file = os.path.join(script_dir, 'parrot.jpg')

# Initialize the TF interpreter
interpreter = edgetpu.make_interpreter(model_file)
interpreter.allocate_tensors()

# Resize the image
size = common.input_size(interpreter)
image = Image.open(image_file).convert('RGB').resize(size, Image.ANTIALIAS)

# Run an inference
common.set_input(interpreter, image)
interpreter.invoke()
classes = classify.get_classes(interpreter, top_k=1)

# Print the result
labels = dataset.read_label_file(label_file)
for c in classes:
  print('%s: %.5f' % (labels.get(c.id, c.id), c.score))

For more examples using the PyCoral API, see our examples on GitHub.

You can also find more example code and projects on our examples page.

For more advanced features using the Coral API, also check out how to run multiple models with multiple Edge TPUs and perform on-device transfer-learning.

Update existing TF Lite code for the Edge TPU

If you already have code using the TensorFlow Lite API, then you essentially need to change just one line of code: When you instantiate the TensorFlow Lite Interpreter, also specify the Edge TPU runtime library as a delegate.

Just follow these steps convert your existing code for the Edge TPU:

1. Install the latest version of the TensorFlow Lite API by following the TensorFlow Lite Python quickstart.

2. In your Python code, import the tflite_runtime module.

   Open the Python file where you'll run inference with the Interpreter API. (For an example, see the TensorFlow Lite code, label_image.py).

   Instead of using import tensorflow as tf, load the tflite_runtime package like this:

   import tflite_runtime.interpreter as tflite

3. Add the Edge TPU delegate when constructing the Interpreter.

   For example, your TensorFlow Lite code will ordinarily have a line like this:

   interpreter = tflite.Interpreter(model_path)

   So change it to this:

   interpreter = tflite.Interpreter(model_path,
     experimental_delegates=[tflite.load_delegate('libedgetpu.so.1')])

   The file passed to load_delegate() is the Edge TPU runtime library, and you should have installed it when you first set up your device. The filename you must use here depends on your host operating system, as follows:

   • Linux: libedgetpu.so.1
   • macOS: libedgetpu.1.dylib
   • Windows: edgetpu.dll

That's it.

Now when you run a model that's compiled for the Edge TPU, TensorFlow Lite delegates the compiled portions of the graph to the Edge TPU.

For a complete example that runs on the Edge TPU, see our image classification example using only TensorFlow Lite.

You can also find more example code and projects on our examples page.