4-6-8-CustomDatasetKWSModel.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pO4-CY_TCZZS"
      },
      "source": [
        "# Training Your Custom Dataset Keyword Spotting Model\n",
        "\n",
        "It is now time for you to train your own custom keyword spotting model using your Custom Dataset!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y0HQ7WkLUxD4"
      },
      "source": [
        "#Setup\n",
        "\n",
        "### Import packages\n",
        "Clone the TensorFlow Github Repository, which contains the relevant code required to run this assignment."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "olCcGuF7GRVO"
      },
      "outputs": [],
      "source": [
        "!wget https://github.com/tensorflow/tensorflow/archive/v2.4.1.zip\n",
        "!unzip v2.4.1.zip &> 0\n",
        "!mv tensorflow-2.4.1/ tensorflow/"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "W2Xne8hbyzT1"
      },
      "source": [
        "### Important - Reverting to TF1 - Start\n",
        "We need to install Tensorflow V1 in order to use the required features for Keyword Spotting. \n",
        "\n",
        "**Run the following cell and click \"restart runtime\" that appears once it completes.**\n",
        "\n",
        "**Do not rerun the previous cells after restarting the runtime.**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "G43nVAIdy3uD"
      },
      "outputs": [],
      "source": [
        "#!pip uninstall tensorflow -y\n",
        "#!pip install tensorflow-gpu==1.15"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yAPKHaKmy5RN"
      },
      "source": [
        "### Important - Reverting to TF1 - End\n",
        "\n",
        "**Make sure to click the \"RESTART RUNTIME\" button in the output of the previous cell. Then move onto the next cell!**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mWIJQmuTy9HU"
      },
      "outputs": [],
      "source": [
        "import tensorflow.compat.v1 as tf\n",
        "# assert(tf.VERSION == \"1.15.0\") #Make sure you install Tensorflow 1.15.0 and restart the runtime or this will fail!\n",
        "# Now that the runtime is set import things!\n",
        "import sys\n",
        "# We add this path so we can import the speech processing modules.\n",
        "sys.path.append(\"/content/tensorflow/tensorflow/examples/speech_commands/\")\n",
        "import input_data\n",
        "import models\n",
        "import numpy as np\n",
        "import glob\n",
        "import os\n",
        "import re\n",
        "import shutil\n",
        "from google.colab import files\n",
        "!pip install ffmpeg-python &> 0\n",
        "!apt-get update && apt-get -qq install xxd"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_QpT5x5sL0RA"
      },
      "source": [
        "### Import your Custom Dataset\n",
        "First we are going to download Pete's dataset to use as a base set of \"other words\" and \"background noise\" that you can build ontop of for your dataset. We have found that doing this will make your model work a lot better, especially if you are training it with a small amount of custom data! We STRONGLY suggest you follow this approach as it will make a large impact on your results!\n",
        "\n",
        "**Note: this may take a couple of minutes to run!**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tE7SthPtL1lY"
      },
      "outputs": [],
      "source": [
        "!wget https://storage.googleapis.com/download.tensorflow.org/data/speech_commands_v0.02.tar.gz\n",
        "DATASET_DIR =  'dataset/'\n",
        "!mkdir dataset\n",
        "!tar -xf speech_commands_v0.02.tar.gz -C 'dataset'\n",
        "!rm -r -f speech_commands_v0.02.tar.gz"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rmqY3_PSMqvl"
      },
      "source": [
        "Now you'll need to upload your all of your custom audio files that you recorded using the Open Speech Recording tool (aka the ```*.ogg``` files). **Note: you can select multiple files and upload them all at once!** \n",
        "\n",
        "If you are having trouble uploading files because your internet bandwidth is too slow feel free to skip this step and you can instead pick from the words in Pete's dataset, just like -- for those of you that took Course 2 -- you did in the KWS assignment.\n",
        "\n",
        "Pete's dataset includes the following words:\n",
        "\n",
        "Options for target words are (PICK FROM THIS LIST FOR BEST RESULTS): \"one\", \"two\", \"three\", \"four\", \"five\", \"six\", \"seven\", \"eight\", \"nine\", \"yes\", \"no\", \"up\", \"down\", \"left\", \"right\", \"on\", \"off\", \"stop\", \"go\", “backward”, “forward”, “follow”, “learn”,\n",
        "\n",
        "Additional words that will be used to help train the \"unkown\" label are: \"bed\", \"bird\", \"cat\", \"dog\", \"happy\", \"house\", \"marvin\", \"sheila\", \"tree\", \"wow\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ppOPH3NfS5W4"
      },
      "outputs": [],
      "source": [
        "uploaded = files.upload()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BU8wwiKmTU8m"
      },
      "source": [
        "Then we can convert them into correctly trimmed WAV files and then store them in the appropriate folders in the ```DATASET_DIR```.\n",
        "We will use Pete's extract_loudest_section tool which you can find more documentation about here: https://github.com/petewarden/extract_loudest_section"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PvD-afuDTLGS"
      },
      "outputs": [],
      "source": [
        "# convert the ogg files to wavs\n",
        "!mkdir wavs\n",
        "!find *.ogg -print0 | xargs -0 basename -s .ogg | xargs -I {} ffmpeg -i {}.ogg -ar 16000 wavs/{}.wav\n",
        "!rm -r -f *.ogg\n",
        "\n",
        "# then use pete's tool to only extract 1 second clips from them for use with the KWS pipeline\n",
        "!mkdir trimmed_wavs\n",
        "!git clone https://github.com/petewarden/extract_loudest_section.git\n",
        "!make -C extract_loudest_section/\n",
        "!/tmp/extract_loudest_section/gen/bin/extract_loudest_section 'wavs/*.wav' trimmed_wavs/\n",
        "!rm -r -f /wavs"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "99ydxbONTgW-"
      },
      "outputs": [],
      "source": [
        "# Store them in the appropriate folders\n",
        "data_index = {}\n",
        "os.chdir('trimmed_wavs')\n",
        "search_path = os.path.join('*.wav')\n",
        "for wav_path in glob.glob(search_path):\n",
        "    original_wav_path = wav_path\n",
        "    parts = wav_path.split('_')\n",
        "    if len(parts) > 2:\n",
        "        wav_path = parts[0] + '_' + ''.join(parts[1:])\n",
        "    matches = re.search('([^/_]+)_([^/_]+)\\.wav', wav_path)\n",
        "    if not matches:\n",
        "        raise Exception('File name not in a recognized form:\"%s\"' % wav_path)\n",
        "    word = matches.group(1).lower()\n",
        "    instance = matches.group(2).lower()\n",
        "    if not word in data_index:\n",
        "      data_index[word] = {}\n",
        "    if instance in data_index[word]:\n",
        "        raise Exception('Audio instance already seen:\"%s\"' % wav_path)\n",
        "    data_index[word][instance] = original_wav_path\n",
        "\n",
        "output_dir = os.path.join('..', 'dataset')\n",
        "try:\n",
        "    os.mkdir(output_dir)\n",
        "except:\n",
        "    pass\n",
        "for word in data_index:\n",
        "  word_dir = os.path.join(output_dir, word)\n",
        "  try:\n",
        "      os.mkdir(word_dir)\n",
        "      print('Created dir: ' + word_dir)\n",
        "  except:\n",
        "      print('Storing in existing dir: ' + word_dir)\n",
        "  for instance in data_index[word]:\n",
        "    wav_path = data_index[word][instance]\n",
        "    output_path = os.path.join(word_dir, instance + '.wav')\n",
        "    shutil.copyfile(wav_path, output_path)\n",
        "os.chdir('..')\n",
        "!rm -r -f trimmed_wavs"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "koCYgUz6U8SW"
      },
      "source": [
        "If you would like to download your full custom dataset as a zipfile for use later you can run the commented out cell below. And then download it from the Files area by clicking on the three dots next to the zip file. Note: this command will take a little while to run as the combination of your data and Pete's dataset will be relatively large!\n",
        "\n",
        "If you would like to create a zip of just your dataset please re-run this colab (or factory reset it to wipe the memory) and skip the line that downloads Pete's dataset. That said, we STRONGLY suggest you do use Pete's dataset for any actual training you do!\n",
        "\n",
        "Finally if you have room in your Google Drive and would like to load and store your dataset from there in the future you can mount your Google Drive in Colab [as described in this blog post](https://towardsdatascience.com/downloading-datasets-into-google-drive-via-google-colab-bcb1b30b0166)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Bg_AHa7VU7wG"
      },
      "outputs": [],
      "source": [
        "#!zip -r myKWSDataset.zip dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XaVtYN4nlCft"
      },
      "source": [
        "### Configure Your Model!\n",
        "Ok now that your custom dataset is all ready to go we'll need to select your keywords and model settings with which to train!\n",
        "\n",
        "```WANTED_WORDS``` = A comma-delimited string of the words you want to train for (e.g., \"yes,no\"). \n",
        "\n",
        "**Make sure to input the keywords you collected!**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ludfxbNIaegy"
      },
      "outputs": [],
      "source": [
        "WANTED_WORDS = \"yes,no\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "P8VBOW8Ob7ys"
      },
      "source": [
        "The number of training steps and learning rates can be specified as comma-separated strings to define the amount/rate at each stage. For example, ```TRAINING_STEPS=\"12000,3000\"``` and ```LEARNING_RATE=\"0.001,0.0001\"``` will run 12,000 training steps with a rate of 0.001 followed by 3,000 final steps with a learning rate of 0.0001. These are good default values to work off of when you choose your values as the course staff has gotten this to work well with those values in the past!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "m1eWOvdfcPlo"
      },
      "outputs": [],
      "source": [
        "TRAINING_STEPS = \"6000,1000\"\n",
        "LEARNING_RATE = \"0.001,0.0001\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0q36xx9vcQVp"
      },
      "source": [
        "We suggest you leave the ```MODEL_ARCHITECTURE``` as tiny_conv the first time but if you would like to do this again and explore additional models some options are: ```single_fc, conv, low_latency_conv, low_latency_svdf, tiny_embedding_conv```. **Do remember if you switch the model type you may need to update the C++ code to include the ```tflite::AllOpsResolver``` to make sure you have all of the neccessary ops!**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "V3Y1phllccMn"
      },
      "outputs": [],
      "source": [
        "MODEL_ARCHITECTURE = 'tiny_conv'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vhj3jQr0cfCt"
      },
      "outputs": [],
      "source": [
        "# Calculate the total number of steps, which is used to identify the checkpoint\n",
        "# file name.\n",
        "TOTAL_STEPS = str(sum(map(lambda string: int(string), TRAINING_STEPS.split(\",\"))))\n",
        "\n",
        "# Print the configuration to confirm it\n",
        "print(\"Training these words: %s\" % WANTED_WORDS)\n",
        "print(\"Training steps in each stage: %s\" % TRAINING_STEPS)\n",
        "print(\"Learning rate in each stage: %s\" % LEARNING_RATE)\n",
        "print(\"Total number of training steps: %s\" % TOTAL_STEPS)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gCgeOpvY9pAi"
      },
      "source": [
        "**We suggest that you do not modify** the following constants as they include filepaths used in this notebook and data that is shared during training and inference."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Nd1iM1o2ymvA"
      },
      "outputs": [],
      "source": [
        "# Calculate the percentage of 'silence' and 'unknown' training samples required\n",
        "# to ensure that we have equal number of samples for each label.\n",
        "number_of_labels = WANTED_WORDS.count(',') + 1\n",
        "number_of_total_labels = number_of_labels + 2 # for 'silence' and 'unknown' label\n",
        "equal_percentage_of_training_samples = int(100.0/(number_of_total_labels))\n",
        "SILENT_PERCENTAGE = equal_percentage_of_training_samples\n",
        "UNKNOWN_PERCENTAGE = equal_percentage_of_training_samples\n",
        "\n",
        "# Constants used during training only\n",
        "VERBOSITY = 'DEBUG'\n",
        "EVAL_STEP_INTERVAL = '1000'\n",
        "SAVE_STEP_INTERVAL = '1000'\n",
        "\n",
        "# Constants for training directories and filepaths\n",
        "LOGS_DIR = 'logs/'\n",
        "TRAIN_DIR = 'train/' # for training checkpoints and other files.\n",
        "\n",
        "# Constants for inference directories and filepaths\n",
        "import os\n",
        "MODELS_DIR = 'models'\n",
        "if not os.path.exists(MODELS_DIR):\n",
        "  os.mkdir(MODELS_DIR)\n",
        "MODEL_TF = os.path.join(MODELS_DIR, 'KWS_custom.pb')\n",
        "MODEL_TFLITE = os.path.join(MODELS_DIR, 'KWS_custom.tflite')\n",
        "FLOAT_MODEL_TFLITE = os.path.join(MODELS_DIR, 'KWS_custom_float.tflite')\n",
        "MODEL_TFLITE_MICRO = os.path.join(MODELS_DIR, 'KWS_custom.cc')\n",
        "SAVED_MODEL = os.path.join(MODELS_DIR, 'KWS_custom_saved_model')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZiKisyFN4R6W"
      },
      "source": [
        "**Be careful if you modify** the following constants as they will have downstream effects on the C++ code which you will then have to change. This mainly relate to hyperparaemeters for quantization and preprocessing. The first time you train a custom model **we suggest you do not modify these as well.**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gbdtOiLV4SMD"
      },
      "outputs": [],
      "source": [
        "# Constants which are shared during training and inference\n",
        "PREPROCESS = 'micro'\n",
        "WINDOW_STRIDE = 20\n",
        "\n",
        "# Constants for Quantization\n",
        "QUANT_INPUT_MIN = 0.0\n",
        "QUANT_INPUT_MAX = 26.0\n",
        "QUANT_INPUT_RANGE = QUANT_INPUT_MAX - QUANT_INPUT_MIN\n",
        "\n",
        "# Constants for audio process during Quantization and Evaluation\n",
        "SAMPLE_RATE = 16000\n",
        "CLIP_DURATION_MS = 1000\n",
        "WINDOW_SIZE_MS = 30.0\n",
        "FEATURE_BIN_COUNT = 40\n",
        "BACKGROUND_FREQUENCY = 0.8\n",
        "BACKGROUND_VOLUME_RANGE = 0.1\n",
        "TIME_SHIFT_MS = 100.0\n",
        "\n",
        "# Use the custom local dataset and set the tes/val/train split\n",
        "DATA_URL = ''\n",
        "VALIDATION_PERCENTAGE = 10\n",
        "TESTING_PERCENTAGE = 10"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UczQKtqLi7OJ"
      },
      "source": [
        "# Train the model\n",
        "\n",
        "### Load in TensorBoard to visulaize the training process.\n",
        "\n",
        "As training progresses you should see the training status show up in the Tensorboard area. If this works it is very helpful for analyzing your training progress. Unfortunately, the staff has found that it sometimes doesn't start showing data for a while (~15 minutes) and sometimes doesn't show data until training completes (and instead shows ```No dashboards are active for the current data set```.). If it is working and then stops updating look to the top of the cell and click reconnect."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ozkJqgSkU_qI"
      },
      "outputs": [],
      "source": [
        "%load_ext tensorboard\n",
        "%tensorboard --logdir {LOGS_DIR}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gweRC9eWVdg2"
      },
      "source": [
        "### Launch Training\n",
        "\n",
        "**ARE YOU USING A GPU INSTANCE!?!?!**\n",
        "Make sure to check ```runtime->change runtime type``` as a GPU runtime will take ~2 hours to complete training whereas a CPU runtime with take ~10 hours!\n",
        "\n",
        "**Importantly, every hour or so you should make sure to remind colab you are still there by clicking around in a cell or clicking reconnect on the tensorboard cell (after 90 minutes of no activity Colab may kill the instance)!** But/and beyond that you can absolutely minimize the window and continue with other work.\n",
        "\n",
        "**If you run training a second time** we suggest restarting your runtime ```runtime->restart runtime``` in order to ensure all caches are cleared. However, do note that doing so will clear all cells and so you will have to run the entire file again.\n",
        "\n",
        "If you would like to get more information on the training script you can find the source code for the script [here](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/examples/speech_commands/train.py). In short it sets up the optimizer and preprocessor based on all of the flags we pass in!\n",
        "\n",
        "Finally, by setting the ```VERBOSITY = 'DEBUG'``` above be aware that the training cell will print A LOT of information. Specifically you will get the accuracy and loss at each step as well as a confusion matrix every 1000 steps. We hope that is helpful in case TensorBoard fails to work. If you would like to run with less printouts you can change the setting to ```WARN``` or ```FATAL```. You will find this in the \"Configure Your Model!\" section."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RZw3VNlnla-J"
      },
      "outputs": [],
      "source": [
        "!python tensorflow/tensorflow/examples/speech_commands/train.py \\\n",
        "  --data_dir={DATASET_DIR} \\\n",
        "  --data_url={DATA_URL} \\\n",
        "  --wanted_words={WANTED_WORDS} \\\n",
        "  --silence_percentage={SILENT_PERCENTAGE} \\\n",
        "  --unknown_percentage={UNKNOWN_PERCENTAGE} \\\n",
        "  --preprocess={PREPROCESS} \\\n",
        "  --window_stride={WINDOW_STRIDE} \\\n",
        "  --model_architecture={MODEL_ARCHITECTURE} \\\n",
        "  --how_many_training_steps={TRAINING_STEPS} \\\n",
        "  --learning_rate={LEARNING_RATE} \\\n",
        "  --train_dir={TRAIN_DIR} \\\n",
        "  --summaries_dir={LOGS_DIR} \\\n",
        "  --verbosity={VERBOSITY} \\\n",
        "  --eval_step_interval={EVAL_STEP_INTERVAL} \\\n",
        "  --save_step_interval={SAVE_STEP_INTERVAL}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pjBn-NapWsho"
      },
      "source": [
        "# Generating your Model\n",
        "Just like with the pre-trained model we will now take the final checkpoint and convert it into a quantized TensorFlow Lite model."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XQUJLrdS-ftl"
      },
      "source": [
        "### Generate a TensorFlow Model for Inference\n",
        "\n",
        "Combine relevant training results (graph, weights, etc) into a single file for inference. This process is known as freezing a model and the resulting model is known as a frozen model/graph, as it cannot be further re-trained after this process."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xyc3_eLh9sAg"
      },
      "outputs": [],
      "source": [
        "!rm -rf {SAVED_MODEL}\n",
        "!python tensorflow/tensorflow/examples/speech_commands/freeze.py \\\n",
        "--wanted_words=$WANTED_WORDS \\\n",
        "--window_stride_ms=$WINDOW_STRIDE \\\n",
        "--preprocess=$PREPROCESS \\\n",
        "--model_architecture=$MODEL_ARCHITECTURE \\\n",
        "--start_checkpoint=$TRAIN_DIR$MODEL_ARCHITECTURE'.ckpt-'{TOTAL_STEPS} \\\n",
        "--save_format=saved_model \\\n",
        "--output_file={SAVED_MODEL}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_DBGDxVI-nKG"
      },
      "source": [
        "### Generate a TensorFlow Lite Model\n",
        "\n",
        "Convert the frozen graph into a TensorFlow Lite model, which is fully quantized for use with embedded devices. The following cell will also print the model size."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rNQdAplJV1fz"
      },
      "outputs": [],
      "source": [
        "model_settings = models.prepare_model_settings(\n",
        "    len(input_data.prepare_words_list(WANTED_WORDS.split(','))),\n",
        "    SAMPLE_RATE, CLIP_DURATION_MS, WINDOW_SIZE_MS,\n",
        "    WINDOW_STRIDE, FEATURE_BIN_COUNT, PREPROCESS)\n",
        "audio_processor = input_data.AudioProcessor(\n",
        "    DATA_URL, DATASET_DIR,\n",
        "    SILENT_PERCENTAGE, UNKNOWN_PERCENTAGE,\n",
        "    WANTED_WORDS.split(','), VALIDATION_PERCENTAGE,\n",
        "    TESTING_PERCENTAGE, model_settings, LOGS_DIR)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "U-rvCCM2G58N"
      },
      "source": [
        "**Note: if the below cell fails it might be because you do not have enough data to have 100 recordings in the representative dataset!** If this happens you will see an error that says something like ```ValueError: cannot reshape array of size 0 into shape (1,1960)```. To help you fix this we have added a ```print(i)``` into the loop. As such, all you have to do is change the ```REP_DATA_SIZE``` variable to be equal to the last integer value printed out by the loop and then re-run the cell!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lBj_AyCh1cC0"
      },
      "outputs": [],
      "source": [
        "REP_DATA_SIZE = 100\n",
        "with tf.Session() as sess:\n",
        "  # with tf.compat.v1.Session() as sess: #replaces the above line for use with TF2.x\n",
        "  float_converter = tf.lite.TFLiteConverter.from_saved_model(SAVED_MODEL)\n",
        "  float_tflite_model = float_converter.convert()\n",
        "  float_tflite_model_size = open(FLOAT_MODEL_TFLITE, \"wb\").write(float_tflite_model)\n",
        "  print(\"Float model is %d bytes\" % float_tflite_model_size)\n",
        "\n",
        "  converter = tf.lite.TFLiteConverter.from_saved_model(SAVED_MODEL)\n",
        "  converter.optimizations = [tf.lite.Optimize.DEFAULT]\n",
        "  converter.inference_input_type = tf.lite.constants.INT8\n",
        "  # converter.inference_input_type = tf.compat.v1.lite.constants.INT8 #replaces the above line for use with TF2.x   \n",
        "  converter.inference_output_type = tf.lite.constants.INT8\n",
        "  # converter.inference_output_type = tf.compat.v1.lite.constants.INT8 #replaces the above line for use with TF2.x\n",
        "  def representative_dataset_gen():\n",
        "    for i in range(REP_DATA_SIZE):\n",
        "      data, _ = audio_processor.get_data(1, i*1, model_settings,\n",
        "                                         BACKGROUND_FREQUENCY, \n",
        "                                         BACKGROUND_VOLUME_RANGE,\n",
        "                                         TIME_SHIFT_MS,\n",
        "                                         'testing',\n",
        "                                         sess)\n",
        "      flattened_data = np.array(data.flatten(), dtype=np.float32).reshape(1, 1960)\n",
        "      print(i)\n",
        "      yield [flattened_data]\n",
        "  converter.representative_dataset = representative_dataset_gen\n",
        "  tflite_model = converter.convert()\n",
        "  tflite_model_size = open(MODEL_TFLITE, \"wb\").write(tflite_model)\n",
        "  print(\"Quantized model is %d bytes\" % tflite_model_size)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EeLiDZTbLkzv"
      },
      "source": [
        "### Testing the accuracy after Quantization\n",
        "\n",
        "Verify that the model we've exported is still accurate, using the TF Lite Python API and our test set."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wQsEteKRLryJ"
      },
      "outputs": [],
      "source": [
        "# Helper function to run inference\n",
        "def run_tflite_inference_testSet(tflite_model_path, model_type=\"Float\"):\n",
        "  #\n",
        "  # Load test data\n",
        "  #\n",
        "  np.random.seed(0) # set random seed for reproducible test results.\n",
        "  with tf.Session() as sess:\n",
        "    # with tf.compat.v1.Session() as sess: #replaces the above line for use with TF2.x\n",
        "    test_data, test_labels = audio_processor.get_data(\n",
        "        -1, 0, model_settings, BACKGROUND_FREQUENCY, BACKGROUND_VOLUME_RANGE,\n",
        "        TIME_SHIFT_MS, 'testing', sess)\n",
        "  test_data = np.expand_dims(test_data, axis=1).astype(np.float32)\n",
        "\n",
        "  #\n",
        "  # Initialize the interpreter\n",
        "  #\n",
        "  interpreter = tf.lite.Interpreter(tflite_model_path)\n",
        "  interpreter.allocate_tensors()\n",
        "  input_details = interpreter.get_input_details()[0]\n",
        "  output_details = interpreter.get_output_details()[0]\n",
        "  \n",
        "  #\n",
        "  # For quantized models, manually quantize the input data from float to integer\n",
        "  #\n",
        "  if model_type == \"Quantized\":\n",
        "    input_scale, input_zero_point = input_details[\"quantization\"]\n",
        "    test_data = test_data / input_scale + input_zero_point\n",
        "    test_data = test_data.astype(input_details[\"dtype\"])\n",
        "\n",
        "  #\n",
        "  # Evaluate the predictions\n",
        "  #\n",
        "  correct_predictions = 0\n",
        "  for i in range(len(test_data)):\n",
        "    interpreter.set_tensor(input_details[\"index\"], test_data[i])\n",
        "    interpreter.invoke()\n",
        "    output = interpreter.get_tensor(output_details[\"index\"])[0]\n",
        "    top_prediction = output.argmax()\n",
        "    correct_predictions += (top_prediction == test_labels[i])\n",
        "\n",
        "  print('%s model accuracy is %f%% (Number of test samples=%d)' % (\n",
        "      model_type, (correct_predictions * 100) / len(test_data), len(test_data)))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "l-pD52Na6jRa"
      },
      "outputs": [],
      "source": [
        "# Compute float model accuracy\n",
        "run_tflite_inference_testSet(FLOAT_MODEL_TFLITE)\n",
        "\n",
        "# Compute quantized model accuracy\n",
        "run_tflite_inference_testSet(MODEL_TFLITE, model_type='Quantized')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nf6CQuooWQy5"
      },
      "source": [
        "### Generate a TensorFlow Lite for Microcontrollers Model\n",
        "To convert the TensorFlow Lite quantized model into a C source file that can be loaded by TensorFlow Lite for Microcontrollers on Arduino we simply need to use the ```xxd``` tool to convert the ```.tflite``` file into a ```.cc``` file (just like we did in the pervious section)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qwStW8XCWXXX"
      },
      "outputs": [],
      "source": [
        "!xxd -i {MODEL_TFLITE} > {MODEL_TFLITE_MICRO}\n",
        "REPLACE_TEXT = MODEL_TFLITE.replace('/', '_').replace('.', '_')\n",
        "!sed -i 's/'{REPLACE_TEXT}'/g_model/g' {MODEL_TFLITE_MICRO}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "P-AhT48A17Eo"
      },
      "source": [
        "That's it! You've successfully converted your TensorFlow Lite model into a TensorFlow Lite for Microcontrollers model! Run the cell below to print out its contents which we'll need for our next step, deploying the model using the Arudino IDE!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "79XIezycHhp6"
      },
      "outputs": [],
      "source": [
        "!cat {MODEL_TFLITE_MICRO}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nQxu8s15HbfJ"
      },
      "source": [
        "To download your model for use at a later date:\n",
        "\n",
        "1. On the left of the UI click on the folder icon\n",
        "2. Click on the three dots to the right of the ```.cc``` file you just generated and select \"download.\" The file can be found at ```models/{MODEL_TFLITE_MICRO}``` which by default is ```models/KWS_custom.cc```\n",
        "\n",
        "Next we'll deploy that model using the Arduino IDE."
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "4-6-8-CustomDatasetKWSModel.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}