Chapter05/ColabNotebooks/prepare_model.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "chapter05.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7xc2grXmGVo4"
      },
      "source": [
        "# Chapter 5 - Indoor Scene Classification with TFLu and the Arduino Nano"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%tensorflow_version 2.x"
      ],
      "metadata": {
        "id": "bAenHrKFofOR"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bCS8r9jvlfFi"
      },
      "source": [
        "### Python libraries"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9IzYsIvDlgx_"
      },
      "source": [
        "import numpy as np\n",
        "import pathlib\n",
        "import tensorflow as tf\n",
        "import tensorflow_datasets as tfds\n",
        "import zipfile\n",
        "from tensorflow.keras.applications.mobilenet_v2 import MobileNetV2\n",
        "from tensorflow.keras.models import Model"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EJchRFEmnwvl"
      },
      "source": [
        "### Constants"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6VGtgqiVFaY1"
      },
      "source": [
        "MODEL_ALPHA = 0.35\n",
        "MODEL_INPUT_WIDTH = 48\n",
        "MODEL_INPUT_HEIGHT = 48\n",
        "TFL_MODEL_FILE = \"indoor_scene_recognition.tflite\"\n",
        "TFL_MODEL_HEADER_FILE = \"indoor_scene_recognition_model.h\"\n",
        "TF_MODEL = \"indoor_scene_recognition\""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Transfer learning with Keras"
      ],
      "metadata": {
        "id": "3J8L3YW_htYw"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Unzip the dataset (dataset.zip)"
      ],
      "metadata": {
        "id": "YSlp-mPQhwuY"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Tp1VN1OZGuEn"
      },
      "source": [
        "with zipfile.ZipFile(\"dataset.zip\", 'r') as zip_ref:\n",
        "    zip_ref.extractall(\".\")\n",
        "data_dir = \"dataset\""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HMqMhiuvGpam"
      },
      "source": [
        "### Prepare the train (80%) and validation (20%) datasets"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jKZ4L-wsBktr"
      },
      "source": [
        "train_ds = tf.keras.utils.image_dataset_from_directory(\n",
        "  data_dir,\n",
        "  validation_split=0.2,\n",
        "  subset=\"training\",\n",
        "  seed=123,\n",
        "  interpolation=\"bilinear\",\n",
        "  image_size=(MODEL_INPUT_WIDTH, MODEL_INPUT_HEIGHT)\n",
        "  )\n",
        "\n",
        "val_ds = tf.keras.utils.image_dataset_from_directory(\n",
        "  data_dir,\n",
        "  validation_split=0.2,\n",
        "  subset=\"validation\",\n",
        "  seed=123,\n",
        "  interpolation=\"bilinear\",\n",
        "  image_size=(MODEL_INPUT_WIDTH, MODEL_INPUT_HEIGHT)\n",
        "  )"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Get the name of the classes"
      ],
      "metadata": {
        "id": "mWxKsXM0iMxj"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zLWkCbjYDVHJ"
      },
      "source": [
        "class_names = train_ds.class_names\n",
        "num_classes = len(class_names)\n",
        "print(class_names)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vEsZQCsUHaeb"
      },
      "source": [
        "### Rescale the pixel values from [0, 255] tp [-1, 1]"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-lFJZj50HhMV"
      },
      "source": [
        "rescale = tf.keras.layers.Rescaling(1./255, offset= -1)\n",
        "train_ds = train_ds.map(lambda x, y: (rescale(x), y))\n",
        "val_ds   = val_ds.map(lambda x, y: (rescale(x), y))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OK0nFgtFSwsF"
      },
      "source": [
        "### Import the MobileNet v2 pre-trained model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DbhSn1Bl-LZ7"
      },
      "source": [
        "# https://github.com/keras-team/keras-applications/blob/master/keras_applications/mobilenet_v2.py\n",
        "base_model = MobileNetV2(input_shape=(MODEL_INPUT_WIDTH, MODEL_INPUT_HEIGHT, 3),\n",
        "                         include_top=False,\n",
        "                         weights='imagenet',\n",
        "                         alpha=0.35)\n",
        "base_model.trainable = False\n",
        "\n",
        "feat_extr = base_model"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lOCL5dx-MmlD"
      },
      "source": [
        "### Augment the input data"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "augmen = tf.keras.Sequential([\n",
        "  tf.keras.layers.experimental.preprocessing.RandomRotation(0.2),\n",
        "  tf.keras.layers.experimental.preprocessing.RandomFlip('horizontal'),\n",
        "])\n",
        "\n",
        "train_ds = train_ds.map(lambda x, y: (augmen(x), y))\n",
        "val_ds   = val_ds.map(lambda x, y: (augmen(x), y))"
      ],
      "metadata": {
        "id": "7ebIppPRjUKF"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Prepare the classification head"
      ],
      "metadata": {
        "id": "Qhw11eldjZoz"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gC5CVjRNMn7q"
      },
      "source": [
        "global_avg_layer = tf.keras.layers.GlobalAveragePooling2D()\n",
        "dense_layer = tf.keras.layers.Dense(num_classes, activation='softmax')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IZYMkd0yM5tb"
      },
      "source": [
        "### Build the model architecture"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "r-DTz_x5M7bg"
      },
      "source": [
        "inputs = tf.keras.Input(shape=(MODEL_INPUT_WIDTH, MODEL_INPUT_HEIGHT, 3))\n",
        "x = global_avg_layer(feat_extr.layers[-1].output)\n",
        "x = tf.keras.layers.Dropout(0.2)(x)\n",
        "outputs = dense_layer(x)\n",
        "model = tf.keras.Model(inputs=feat_extr.inputs, outputs=outputs)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Compile the model with a 0.0005 learning rate"
      ],
      "metadata": {
        "id": "X4AoQHO1j1a3"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZHPUqQNijqE5"
      },
      "source": [
        "lr = 0.0005\n",
        "model.compile(\n",
        "  optimizer=tf.keras.optimizers.Adam(learning_rate=lr),\n",
        "  loss=tf.losses.SparseCategoricalCrossentropy(from_logits=False),\n",
        "  metrics=['accuracy'])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Model summary"
      ],
      "metadata": {
        "id": "x4M-8ftzj89-"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WFXJSd6fj415"
      },
      "source": [
        "model.summary()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Train the model with 10 epochs"
      ],
      "metadata": {
        "id": "25W5CUAVj_I9"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "na0GBGTQkavr"
      },
      "source": [
        "model.fit(\n",
        "  train_ds,\n",
        "  validation_data=val_ds,\n",
        "  epochs=10\n",
        ")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Save the TensorFlow model"
      ],
      "metadata": {
        "id": "U0tsAdGikFa_"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rOUI3vtLtYKh"
      },
      "source": [
        "model.save(TF_MODEL)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KntdZZ5ruM7Z"
      },
      "source": [
        "## Preparing and testing the quantized TFLite model"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Unzip the test dataset (test_samples.zip)"
      ],
      "metadata": {
        "id": "srHkdaf_kk97"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "with zipfile.ZipFile(\"test_samples.zip\", 'r') as zip_ref:\n",
        "    zip_ref.extractall(\".\")\n",
        "test_dir = \"dataset\""
      ],
      "metadata": {
        "id": "ZxnVbqxkkspY"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Rescale the pixel values from [0, 255] to [-1, 1]"
      ],
      "metadata": {
        "id": "jzTiU6qQk_Bv"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "test_ds = tf.keras.utils.image_dataset_from_directory(test_dir,\n",
        "                                                      interpolation=\"bilinear\",\n",
        "                                                      image_size=(MODEL_INPUT_WIDTH, MODEL_INPUT_HEIGHT))\n",
        "test_ds  = test_ds.map(lambda x, y: (rescale(x), y))"
      ],
      "metadata": {
        "id": "vYHpua21lGXp"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Quantize the TensorFlow model with the TFLite converter"
      ],
      "metadata": {
        "id": "riqPFiWllkdg"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ALWbwZKBuNl8"
      },
      "source": [
        "repr_ds = test_ds.unbatch()\n",
        "\n",
        "def representative_data_gen():\n",
        "  for i_value, o_value in repr_ds.batch(1).take(48):\n",
        "    yield [i_value]\n",
        "\n",
        "converter = tf.lite.TFLiteConverter.from_saved_model(TF_MODEL)\n",
        "converter.representative_dataset = tf.lite.RepresentativeDataset(representative_data_gen)\n",
        "converter.optimizations = [tf.lite.Optimize.DEFAULT]\n",
        "converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]\n",
        "converter.inference_input_type = tf.int8\n",
        "\n",
        "tfl_model = converter.convert()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Get the TFLite model size in bytes"
      ],
      "metadata": {
        "id": "nQg4FbwzluB8"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "size_tfl_model = len(tfl_model)\n",
        "print(len(tfl_model), \"bytes\")"
      ],
      "metadata": {
        "id": "IFIpV1uUlxiZ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Initialize the TFLite interpreter"
      ],
      "metadata": {
        "id": "1Br6hq8Cl4nU"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ACsaQjxzP2Wd"
      },
      "source": [
        "# Initialize the TFLite interpreter\n",
        "interpreter = tf.lite.Interpreter(model_content=tfl_model)\n",
        "\n",
        "# Allocate the tensors\n",
        "interpreter.allocate_tensors()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Get input quantization parameters"
      ],
      "metadata": {
        "id": "OxxhFk-7l9J2"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Get input/output layer information\n",
        "i_details = interpreter.get_input_details()[0]\n",
        "o_details = interpreter.get_output_details()[0]\n",
        "\n",
        "# Get input quantization parameters.\n",
        "i_quant = i_details[\"quantization_parameters\"]\n",
        "i_scale      = i_quant['scales'][0]\n",
        "i_zero_point = i_quant['zero_points'][0]"
      ],
      "metadata": {
        "id": "nrH3uM8Gl_Hg"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Evaluate the accuracy of the quantized TFLite model"
      ],
      "metadata": {
        "id": "WQcFhKfumT1B"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "test_ds0 = val_ds.unbatch()\n",
        "\n",
        "num_correct_samples = 0\n",
        "num_total_samples   = len(list(test_ds0.batch(1)))\n",
        "\n",
        "for i_value, o_value in test_ds0.batch(1):\n",
        "  i_value = (i_value / i_scale) + i_zero_point\n",
        "  i_value = tf.cast(i_value, dtype=tf.int8)\n",
        "  interpreter.set_tensor(i_details[\"index\"], i_value)\n",
        "  interpreter.invoke()\n",
        "  o_pred = interpreter.get_tensor(o_details[\"index\"])[0]\n",
        "\n",
        "  if np.argmax(o_pred) == o_value:\n",
        "    num_correct_samples += 1\n",
        "\n",
        "print(\"Accuracy:\", num_correct_samples/num_total_samples)"
      ],
      "metadata": {
        "id": "1vlIBfMHmWT-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Convert the TFLite model to C-byte array with xxd"
      ],
      "metadata": {
        "id": "tM-ZaNUHmkBm"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3kAnQgmI0QK_"
      },
      "source": [
        "open(\"model.tflite\", \"wb\").write(tfl_model)\n",
        "!apt-get update && apt-get -qq install xxd\n",
        "!xxd -c 60 -i model.tflite > indoor_scene_recognition.h"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}