Train a Simple Audio Recognition model for microcontroller use

This notebook demonstrates how to train a 20kb Simple Audio Recognition model for TensorFlow Lite for Microcontrollers. It will produce the same model used in the micro_speech example application.

The model is designed to be used with .

View source on GitHub

The notebook runs Python scripts to train and freeze the model, and uses the TensorFlow Lite converter to convert it for use with TensorFlow Lite for Microcontrollers.

Training is much faster using GPU acceleration. Before you proceed, ensure you are using a GPU runtime by going to Runtime -> Change runtime type and selecting GPU. Training 18,000 iterations will take 1.5-2 hours on a GPU runtime.

Configure training

The following os.environ lines can be customized to set the words that will be trained for, and the steps and learning rate of the training. The default values will result in the same model that is used in the micro_speech example. Run the cell to set the configuration:



In [0]:

    
import os

# A comma-delimited list of the words you want to train for.
# The options are: yes,no,up,down,left,right,on,off,stop,go
# All other words will be used to train an "unknown" category.
os.environ["WANTED_WORDS"] = "yes,no"

# The number of steps and learning rates can be specified as comma-separated
# lists to define the rate at each stage. For example,
# TRAINING_STEPS=15000,3000 and LEARNING_RATE=0.001,0.0001
# will run 18,000 training loops in total, with a rate of 0.001 for the first
# 15,000, and 0.0001 for the final 3,000.
os.environ["TRAINING_STEPS"]="15000,3000"
os.environ["LEARNING_RATE"]="0.001,0.0001"

# Calculate the total number of steps, which is used to identify the checkpoint
# file name.
total_steps = sum(map(lambda string: int(string),
                  os.environ["TRAINING_STEPS"].split(",")))
os.environ["TOTAL_STEPS"] = str(total_steps)

# Print the configuration to confirm it
!echo "Training these words: ${WANTED_WORDS}"
!echo "Training steps in each stage: ${TRAINING_STEPS}"
!echo "Learning rate in each stage: ${LEARNING_RATE}"
!echo "Total number of training steps: ${TOTAL_STEPS}"

Install dependencies

Next, we'll install a GPU build of TensorFlow, so we can use GPU acceleration for training.



In [0]:

    
# Replace Colab's default TensorFlow install with a more recent
# build that contains the operations that are needed for training
!pip uninstall -y tensorflow tensorflow_estimator
!pip install -q tf-estimator-nightly==1.14.0.dev2019072901 tf-nightly-gpu==1.15.0.dev20190729

We'll also clone the TensorFlow repository, which contains the scripts that train and freeze the model.



In [0]:

    
# Clone the repository from GitHub
!git clone -q https://github.com/tensorflow/tensorflow
# Check out a commit that has been tested to work
# with the build of TensorFlow we're using
!git -c advice.detachedHead=false -C tensorflow checkout 17ce384df70

Load TensorBoard

Now, set up TensorBoard so that we can graph our accuracy and loss as training proceeds.



In [0]:

    
# Delete any old logs from previous runs
!rm -rf /content/retrain_logs
# Load TensorBoard
%load_ext tensorboard
%tensorboard --logdir /content/retrain_logs

Begin training

Next, run the following script to begin training. The script will first download the training data:



In [0]:

    
!python tensorflow/tensorflow/examples/speech_commands/train.py \
--model_architecture=tiny_conv --window_stride=20 --preprocess=micro \
--wanted_words=${WANTED_WORDS} --silence_percentage=25 --unknown_percentage=25 \
--quantize=1 --verbosity=WARN --how_many_training_steps=${TRAINING_STEPS} \
--learning_rate=${LEARNING_RATE} --summaries_dir=/content/retrain_logs \
--data_dir=/content/speech_dataset --train_dir=/content/speech_commands_train \

Freeze the graph

Once training is complete, run the following cell to freeze the graph.



In [0]:

    
!python tensorflow/tensorflow/examples/speech_commands/freeze.py \
--model_architecture=tiny_conv --window_stride=20 --preprocess=micro \
--wanted_words=${WANTED_WORDS} --quantize=1 --output_file=/content/tiny_conv.pb \
--start_checkpoint=/content/speech_commands_train/tiny_conv.ckpt-${TOTAL_STEPS}

Convert the model

Run this cell to use the TensorFlow Lite converter to convert the frozen graph into the TensorFlow Lite format, fully quantized for use with embedded devices.



In [0]:

    
!toco \
--graph_def_file=/content/tiny_conv.pb --output_file=/content/tiny_conv.tflite \
--input_shapes=1,49,40,1 --input_arrays=Reshape_2 --output_arrays='labels_softmax' \
--inference_type=QUANTIZED_UINT8 --mean_values=0 --std_dev_values=9.8077

The following cell will print the model size, which will be under 20 kilobytes.



In [0]:

    
import os
model_size = os.path.getsize("/content/tiny_conv.tflite")
print("Model is %d bytes" % model_size)

Finally, we use xxd to transform the model into a source file that can be included in a C++ project and loaded by TensorFlow Lite for Microcontrollers.



In [0]:

    
# Install xxd if it is not available
!apt-get -qq install xxd
# Save the file as a C source file
!xxd -i /content/tiny_conv.tflite > /content/tiny_conv.cc
# Print the source file
!cat /content/tiny_conv.cc