Copyright 2020 Google LLC

Licensed under the Apache License, Version 2.0 (the "License")



In [0]:

    
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

Retrain a classification model for Edge TPU with quant-aware training (TF 1.11)

This notebook uses a set of TensorFlow training scripts to perform transfer-learning on a quantization-aware classification model and then convert it for compatibility with the Edge TPU.

Specifically, this tutorial shows you how to perform fine-tuning on the MobileNet V1 model so it can recognize a new set of classes (five types of flowers), using TensorFlow r1.11.

Beware that, compared to a desktop computer, this training can take much longer in Colab because Colab provides limited resources for long-running operations. So you'll likely see faster training speeds if you connect this notebook to a local runtime, or instead follow the tutorial to run this training in Docker (which includes more documentation about this process).

Note: As an alternative approach, check out this other Colab notebook that retrains a classification model for the Edge TPU using post-training quantization (instead of quantization-aware training).

Import TensorFlow



In [0]:

    
! pip uninstall tensorflow -y
! pip install tensorflow==1.11



In [0]:

    
import tensorflow as tf
print(tf.__version__)

Clone the model and training repos



In [0]:

    
! git clone https://github.com/tensorflow/models.git



In [0]:

    
! cd models && git checkout f788046ca876a8820e05b0b48c1fc2e16b0955bc



In [0]:

    
! git clone https://github.com/google-coral/tutorials.git

! cp tutorials/docker/classification/scripts/* models/research/slim/

Convert training data to TFRecord



In [0]:

    
%cd models/research/slim



In [0]:

    
! ./prepare_checkpoint_and_dataset.sh --network_type mobilenet_v1

Perform transfer-learning

The following script takes a couple hours to finish in Colab. (You can shorten by reducing the steps, but that reduces the final accuracy.)

If you didn't already select "Run all" then you should run all remaining cells now. That will ensure the rest of the notebook completes while you are away, avoiding the chance that the Colab runtime times-out and you lose the training data before you download the model.



In [0]:

    
%env NUM_TRAINING_STEPS=300



In [0]:

    
! ./start_training.sh --network_type mobilenet_v1 --train_steps $NUM_TRAINING_STEPS

As training progresses, you can see new checkpoint files appear in the models/research/slim/transfer_learn/train/ directory.



In [0]:

    
! ./run_evaluation.sh --network_type mobilenet_v1

Compile for the Edge TPU



In [0]:

    
! ./convert_checkpoint_to_edgetpu_tflite.sh --network_type mobilenet_v1 --checkpoint_num $NUM_TRAINING_STEPS



In [0]:

    
! curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key add -

! echo "deb https://packages.cloud.google.com/apt coral-edgetpu-stable main" | sudo tee /etc/apt/sources.list.d/coral-edgetpu.list

! sudo apt-get update

! sudo apt-get install edgetpu-compiler



In [0]:

    
%cd transfer_learn/models/

! ls



In [0]:

    
! edgetpu_compiler output_tflite_graph.tflite

Download the files:



In [0]:

    
from google.colab import files

files.download('output_tflite_graph_edgetpu.tflite')
files.download('labels.txt')

If you get a "Failed to fetch" error here, it's probably because the files weren't done saving. So just wait a moment and try again.

Also look out for a browser popup that might need approval to download the files.

Run the model on the Edge TPU

You can now run the model on your Coral device with acceleration on the Edge TPU.

To get started, try using this code for image classification with the TensorFlow Lite API. Just follow the instructions on that page to set up your device, copy the output_tflite_graph_edgetpu.tflite and labels.txt files to your Coral Dev Board or device with a Coral Accelerator, and pass it a flower photo like this:

python3 classify_image.py \
  --model output_tflite_graph_edgetpu.tflite \
  --labels labels.txt \
  --input flower.jpg

Check out more examples for running inference at coral.ai/examples.

Implementation details

All the scripts used in this notebook come from the following locations:

More explanation of the steps in this tutorial is available at https://coral.ai/docs/edgetpu/retrain-classification/.