4-4-8-Flatbuffers.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S9eLyYdy9t4I"
      },
      "source": [
        "# TensorFlow Lite Flatbuffer Manipulation Example\n",
        "\n",
        "It's possible to read, modify, and write TensorFlow Lite model files from Python. This notebook shows you how.\n",
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/tinyMLx/colabs/blob/master/4-4-8-Flatbuffers.ipynb\n",
        "\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/tinyMLx/colabs/blob/master/4-4-8-Flatbuffers.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LN0slo1N-dMT"
      },
      "source": [
        "## Software installation\n",
        "\n",
        "The goal is to build a set of Python classes that represent the types stored inside the TFL Flatbuffer files. To do this we need several dependencies:\n",
        " - The 'flatc' compiler, that converts the model format stored in a text schema to Python accessor classes.\n",
        " - The text schema itself, describing the model format.\n",
        " - The Flatbuffer Python library that the accessor classes rely on.\n",
        "\n",
        "The 'flatc' compiler isn't available as a binary download, so we need to build it from source. The version has to match the Flatbuffer Python library on the system, or the resulting generated code won't work, since it will be trying to use a different API. The FB Python library is at version 1.12.0, so we make sure we download the tagged snapshot of the source at that same version from GitHub."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Gvc8Gv806odl"
      },
      "source": [
        "### Install Flatbuffer Python Library\n",
        "\n",
        "We should already have version 1.12.0 by default on this Colab, but use pip to ensure it's installed, and then import it so it's available from Python."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "v9eegi_vtxW4"
      },
      "outputs": [],
      "source": [
        "!pip install flatbuffers==23.5.26\n",
        "import flatbuffers"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "njMMlz3L69Re"
      },
      "source": [
        "### Build the 'flatc' Compiler\n",
        "\n",
        "We need 'flatc' to generate the Python accessor classes we'll use to read and write the serialized files, and since it's not easily available as a binary, we grab the source code from the right version and build it directly. **This will take a few minutes.**\n",
        "\n",
        "Once the 'flatc' binary has been built, copy it to the `/usr/local/bin` folder so that it's easily accessible as a command."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HiM0ZsxO6NuX"
      },
      "outputs": [],
      "source": [
        "# Build and install the Flatbuffer compiler.\n",
        "%cd /content/\n",
        "!rm -rf flatbuffers*\n",
        "!curl -L \"https://github.com/google/flatbuffers/archive/refs/tags/v23.5.26.zip\" -o flatbuffers.zip\n",
        "!unzip -q flatbuffers.zip\n",
        "!mv flatbuffers-23.5.26 flatbuffers\n",
        "%cd flatbuffers\n",
        "!cmake -G \"Unix Makefiles\" -DCMAKE_BUILD_TYPE=Release\n",
        "!make -j 8\n",
        "!cp flatc /usr/local/bin/"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "D5vlh6BCM9i0"
      },
      "source": [
        "### Fetch the Schema\n",
        "\n",
        "The schema is [a text file that describes the layout of the model file format](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/schema/schema_v3.fbs), and it's part of the TensorFlow source code, so we need to fetch the latest version from GitHub."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4boI3wM00PnS"
      },
      "outputs": [],
      "source": [
        "%cd /content/\n",
        "!rm -rf tensorflow\n",
        "%tensorflow_version 2.x\n",
        "!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": "GS1eAfwRNfvG"
      },
      "source": [
        "### Generate the Python Accessor Classes\n",
        "\n",
        "The 'flatc' compiler takes the information from the text schema, and creates Python code to read and write the data held inside the serialized Flatbuffer file. The Python classes are written into the `tflite` folder, with names like `Model.py`, and define classes like `ModelT` that contain members that you can use to read or write the contents of the data structures defined by the schema."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Xl0_MIlMM6Es"
      },
      "outputs": [],
      "source": [
        "!flatc --python --gen-object-api tensorflow/tensorflow/lite/schema/schema_v3.fbs"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8p_5vLNQ_sFF"
      },
      "source": [
        "### Model Loading and Saving\n",
        "\n",
        "These are some simple wrapper functions that demonstrate how to load data from a file, turn it into a `ModelT` Python object that can be modified, and then save it out to a new file."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "acpsc72lvKRW"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "import tflite.Model as Model\n",
        "\n",
        "def load_model_from_file(model_filename):\n",
        "  with open(model_filename, \"rb\") as file:\n",
        "    buffer_data = file.read()\n",
        "  model_obj = Model.Model.GetRootAsModel(buffer_data, 0)\n",
        "  model = Model.ModelT.InitFromObj(model_obj)\n",
        "  return model\n",
        "\n",
        "def save_model_to_file(model, model_filename):\n",
        "  builder = flatbuffers.Builder(1024)\n",
        "  model_offset = model.Pack(builder)\n",
        "  builder.Finish(model_offset, file_identifier=b'TFL3')\n",
        "  model_data = builder.Output()\n",
        "  with open(model_filename, 'wb') as out_file:\n",
        "    out_file.write(model_data)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vtNYC984AILB"
      },
      "source": [
        "## Download an Example Model\n",
        "\n",
        "This pulls down a trained model from the speech commands tutorial that we can use to test the Flatbuffer loading code."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "56u94mxKsCEM"
      },
      "outputs": [],
      "source": [
        "!curl -O 'https://storage.googleapis.com/download.tensorflow.org/models/tflite/micro/speech_commands_model_2020_04_27.zip'\n",
        "!unzip speech_commands_model_2020_04_27.zip"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "B5SJTcrRBy1X"
      },
      "source": [
        "## Load, Modify, and Save a Model\n",
        "\n",
        "The code below loads the float model, applies a small change to the float weights, and saves it out again. In this case we're just changing the contents of the weights, but any of the other properties of the model can be added, removed, or modified, including tensors, ops, and metadata.\n",
        "\n",
        "I've found the easiest way to understand the syntax used and data available is to look at the generated classes in `tflite/` and [the text schema that describes the format](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/schema/schema_v3.fbs)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zex9zZo01lM4"
      },
      "outputs": [],
      "source": [
        "# Use numpy to make the manipulation of the weight arrays easier.\n",
        "import numpy as np\n",
        "\n",
        "# Load the float model we downloaded as a ModelT object.\n",
        "model = load_model_from_file('/content/speech_commands_model/speech_commands_model_float.tflite')\n",
        "\n",
        "# Loop through all the weights held in the model.\n",
        "for buffer in model.buffers:\n",
        "  # Skip missing or small weight arrays.\n",
        "  if buffer.data is not None and len(buffer.data) > 1024:\n",
        "    # Pull the weight array from the model, and cast it to 32-bit floats since\n",
        "    # we know that's the type for all the weights in this model. In a real\n",
        "    # application we'd need to check the data type from the tensor information\n",
        "    # stored in the model.subgraphs\n",
        "    original_weights = np.frombuffer(buffer.data, dtype=np.float32)\n",
        "\n",
        "    # This is the line where the weights are altered.\n",
        "    # Try replacing it with your own version, for example:\n",
        "    # munged_weights = np.add(original_weights, 0.002)\n",
        "    munged_weights = np.round(original_weights * (1/0.02)) * 0.02\n",
        "\n",
        "    # Write the altered data back into the model.\n",
        "    buffer.data = munged_weights.tobytes()\n",
        "\n",
        "# Save the ModelT object as a TFL Flatbuffer file.\n",
        "save_model_to_file(model, '/content/speech_commands_model/speech_commands_model_modified.tflite')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "C8EwrywLX9gs"
      },
      "source": [
        "## Evaluating the impact of your changes"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_EAgJJNmZL8T"
      },
      "source": [
        "### Setup KWS infrastructure\n",
        "\n",
        "To evaluate the impact of your changes on the accuracy of the speech model we need to load a test data set and some utility classes to read the files and convert them into the right input form for the network.\n",
        "\n",
        "**The full dataset is several gigabytes in size, so it may take a few minutes to download.**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nlrzIA__X9v3"
      },
      "outputs": [],
      "source": [
        "sys.path.append(\"/content/tensorflow/tensorflow/examples/speech_commands/\")\n",
        "import input_data\n",
        "import models\n",
        "\n",
        "# A comma-delimited list of the words you want to train for.\n",
        "# The options are: yes,no,up,down,left,right,on,off,stop,go\n",
        "# All the other words will be used to train an \"unknown\" label and silent\n",
        "# audio data with no spoken words will be used to train a \"silence\" label.\n",
        "WANTED_WORDS = \"yes,no\"\n",
        "\n",
        "# The number of steps and learning rates can be specified as comma-separated\n",
        "# lists to define the rate at each stage. For example,\n",
        "# TRAINING_STEPS=12000,3000 and LEARNING_RATE=0.001,0.0001\n",
        "# will run 12,000 training loops in total, with a rate of 0.001 for the first\n",
        "# 8,000, and 0.0001 for the final 3,000.\n",
        "TRAINING_STEPS = \"12000,3000\"\n",
        "LEARNING_RATE = \"0.001,0.0001\"\n",
        "\n",
        "# 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",
        "# 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 which are shared during training and inference\n",
        "PREPROCESS = 'micro'\n",
        "WINDOW_STRIDE =20\n",
        "MODEL_ARCHITECTURE = 'tiny_conv' # Other options include: single_fc, conv,\n",
        "                      # low_latency_conv, low_latency_svdf, tiny_embedding_conv\n",
        "\n",
        "# Constants used during training only\n",
        "VERBOSITY = 'WARN'\n",
        "EVAL_STEP_INTERVAL = '1000'\n",
        "SAVE_STEP_INTERVAL = '1000'\n",
        "\n",
        "# Constants for training directories and filepaths\n",
        "DATASET_DIR =  'dataset/'\n",
        "LOGS_DIR = 'logs/'\n",
        "TRAIN_DIR = 'train/' # for training checkpoints and other files.\n",
        "\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",
        "DATA_URL = 'https://storage.googleapis.com/download.tensorflow.org/data/speech_commands_v0.02.tar.gz'\n",
        "VALIDATION_PERCENTAGE = 10\n",
        "TESTING_PERCENTAGE = 10\n",
        "\n",
        "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": "code",
      "execution_count": null,
      "metadata": {
        "id": "bMsyppDEYaYP"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "# define our helper function to test the model accuracy\n",
        "def test_model_accuracy(model_filename):\n",
        "  with tf.compat.v1.Session() as sess:\n",
        "    test_data, test_labels = audio_processor.get_data(\n",
        "        -1, 0, model_settings, 0, 0,\n",
        "        0, 'testing', sess)\n",
        "\n",
        "  interpreter = tf.lite.Interpreter(model_filename)\n",
        "  interpreter.allocate_tensors()\n",
        "\n",
        "  input_index = interpreter.get_input_details()[0][\"index\"]\n",
        "\n",
        "  output_index = interpreter.get_output_details()[0][\"index\"]\n",
        "  model_output = interpreter.tensor(output_index)\n",
        "\n",
        "  correct_predictions = 0\n",
        "  for i in range(len(test_data)):\n",
        "    current_input = test_data[i]\n",
        "    current_label = test_labels[i]\n",
        "    flattened_input = np.array(current_input.flatten(), dtype=np.float32).reshape(1, 1960)\n",
        "    interpreter.set_tensor(input_index, flattened_input)\n",
        "    interpreter.invoke()\n",
        "    top_prediction = model_output()[0].argmax()\n",
        "    if top_prediction == current_label:\n",
        "      correct_predictions += 1\n",
        "\n",
        "  print('Accuracy is %f%% (N=%d)' % ((correct_predictions * 100) / len(test_data), len(test_data)))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PLSPYyKxYoNT"
      },
      "source": [
        "### Test your Models\n",
        "\n",
        "Finally we can test both the standard and your modified models!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PDkPbCi_ZYJL"
      },
      "source": [
        "**Float model evaluation**\n",
        "\n",
        "You should see ~91% accuracy with a 67KB model. Note: the exact accuracy may vary by a percent or two as it tests the model on a random sampling of the dataset and therefore sometimes gets particularly lucky or unlucky!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "toSX48DGYjqP"
      },
      "outputs": [],
      "source": [
        "!ls -lah /content/speech_commands_model/speech_commands_model_float.tflite\n",
        "test_model_accuracy('/content/speech_commands_model/speech_commands_model_float.tflite')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RF3T18DWZe7J"
      },
      "source": [
        "**Modified model evaluation**\n",
        "\n",
        "Test the impact of your changes!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UIfNzw7kZqga"
      },
      "outputs": [],
      "source": [
        "!ls -lah /content/speech_commands_model/speech_commands_model_modified.tflite\n",
        "test_model_accuracy('/content/speech_commands_model/speech_commands_model_modified.tflite')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Vzw4-n8HAgtT"
      },
      "source": [
        "If you'd like to try more modified models just scroll back up to the **Load, Modify, and Save a Model** section and re-run the modified model generation step and then skip back down to the above to re-test!"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}