Mnist_Quantisation_aware_training.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Mnist_Quantization_aware_training.ipynb",
      "provenance": [],
      "authorship_tag": "ABX9TyOg2W9Jny6ZpLzKfxSP5LJ8",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/jvishnuvardhan/TF_Lite/blob/master/Mnist_Quantisation_aware_training.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "O5AQ9wJGulrn",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "272d3c2a-631a-4d13-9d5b-d2e9e896076c"
      },
      "source": [
        "! pip uninstall -y tensorflow\n",
        "! pip install -q tf-nightly\n",
        "! pip install -q tensorflow-model-optimization"
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "\u001b[33mWARNING: Skipping tensorflow as it is not installed.\u001b[0m\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EBIDe6ylurYi",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import tempfile\n",
        "import os\n",
        "\n",
        "import tensorflow as tf\n",
        "\n",
        "from tensorflow import keras"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Q2TFprfzu3b7",
        "colab_type": "text"
      },
      "source": [
        "1. **Train a model for MNIST without quantization aware training**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bFVnzYk6uxWx",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 51
        },
        "outputId": "c5d3477e-dc18-419d-dfe2-3e9fd0745693"
      },
      "source": [
        "# Load MNIST dataset\n",
        "mnist = keras.datasets.mnist\n",
        "(train_images, train_labels), (test_images, test_labels) = mnist.load_data()\n",
        "\n",
        "# Normalize the input image so that each pixel value is between 0 to 1.\n",
        "train_images = train_images / 255.0\n",
        "test_images = test_images / 255.0\n",
        "\n",
        "# Define the model architecture.\n",
        "model = keras.Sequential([\n",
        "  keras.layers.InputLayer(input_shape=(28, 28)),\n",
        "  keras.layers.Reshape(target_shape=(28, 28, 1)),\n",
        "  keras.layers.Conv2D(filters=12, kernel_size=(3, 3), activation=tf.nn.relu),\n",
        "  keras.layers.MaxPooling2D(pool_size=(2, 2)),\n",
        "  keras.layers.Flatten(),\n",
        "  keras.layers.Dense(10, activation=tf.nn.softmax)\n",
        "])\n",
        "\n",
        "# Train the digit classification model\n",
        "model.compile(optimizer='adam',\n",
        "              loss='sparse_categorical_crossentropy',\n",
        "              metrics=['accuracy'])\n",
        "\n",
        "model.fit(\n",
        "  train_images,\n",
        "  train_labels,\n",
        "  epochs=1,\n",
        "  validation_split=0.1,\n",
        ")"
      ],
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "1688/1688 [==============================] - 18s 11ms/step - loss: 0.2753 - accuracy: 0.9229 - val_loss: 0.0985 - val_accuracy: 0.9773\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<tensorflow.python.keras.callbacks.History at 0x7f0d3e5c3748>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 3
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "W9ikfE8Ivj1d",
        "colab_type": "text"
      },
      "source": [
        "**2. Clone and fine-tune pre-trained model with quantization aware training**\n",
        "\n",
        "*   apply quantization aware training to the whole model\n",
        "*   All layers are now prefixed by \"quant\"\n",
        "*   resulting model is quantization aware but not quantized (e.g. the weights are float32 instead of int8)\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9CiSoS1xvier",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 51
        },
        "outputId": "f33c8ce2-3948-42af-f9d7-9e0c0576edb1"
      },
      "source": [
        "import tensorflow_model_optimization as tfmot\n",
        "\n",
        "quantize_model = tfmot.quantization.keras.quantize_model\n",
        "\n",
        "# q_aware stands for for quantization aware.\n",
        "q_aware_model = quantize_model(model)\n",
        "\n",
        "# `quantize_model` requires a recompile.\n",
        "q_aware_model.compile(optimizer='adam',\n",
        "              loss='sparse_categorical_crossentropy',\n",
        "              metrics=['accuracy'])\n",
        "\n",
        "#q_aware_model.summary()\n",
        "\n",
        "train_images_subset = train_images[0:1000] # out of 60000\n",
        "train_labels_subset = train_labels[0:1000]\n",
        "\n",
        "q_aware_model.fit(train_images_subset, train_labels_subset,\n",
        "                  batch_size=500, epochs=1, validation_split=0.1)"
      ],
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "2/2 [==============================] - 0s 141ms/step - loss: 0.1245 - accuracy: 0.9633 - val_loss: 0.1346 - val_accuracy: 0.9600\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<tensorflow.python.keras.callbacks.History at 0x7f0d3b6317b8>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 4
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BCL44ao0w7rz",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 51
        },
        "outputId": "19386a22-4ce9-42af-e81d-503889ca9e08"
      },
      "source": [
        "# checking the loss in accuracy between normal model and qaware model\n",
        "_, baseline_model_accuracy = model.evaluate(\n",
        "    test_images, test_labels, verbose=0)\n",
        "\n",
        "_, q_aware_model_accuracy = q_aware_model.evaluate(\n",
        "   test_images, test_labels, verbose=0)\n",
        "\n",
        "print('Baseline test accuracy:', baseline_model_accuracy)\n",
        "print('Quant test accuracy:', q_aware_model_accuracy)"
      ],
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Baseline test accuracy: 0.9689000248908997\n",
            "Quant test accuracy: 0.9682000279426575\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HmXP3wl0yJ1D",
        "colab_type": "text"
      },
      "source": [
        "**3. Create quantized model for TFLite backend**   \n",
        "After this, you have an actually quantized model with int8 weights and uint8 activations."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LbeI8xutyFAB",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 156
        },
        "outputId": "6cb41fb8-bcb2-4f48-b687-c54668725438"
      },
      "source": [
        "converter = tf.lite.TFLiteConverter.from_keras_model(q_aware_model)\n",
        "converter.optimizations = [tf.lite.Optimize.DEFAULT]\n",
        "\n",
        "quantized_tflite_model = converter.convert()"
      ],
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/backend.py:465: set_learning_phase (from tensorflow.python.keras.backend) is deprecated and will be removed after 2020-10-11.\n",
            "Instructions for updating:\n",
            "Simply pass a True/False value to the `training` argument of the `__call__` method of your layer or model.\n",
            "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/training/tracking/tracking.py:105: Model.state_updates (from tensorflow.python.keras.engine.training) is deprecated and will be removed in a future version.\n",
            "Instructions for updating:\n",
            "This property should not be used in TensorFlow 2.0, as updates are applied automatically.\n",
            "INFO:tensorflow:Assets written to: /tmp/tmpjvd7obgu/assets\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9fTwdjvQ2etZ",
        "colab_type": "text"
      },
      "source": [
        "**4. See persistence of accuracy from TF to TFLite**.  \n",
        "Define a helper function to evaluate the TF Lite model on the test dataset."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "makEGxEGx18H",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import numpy as np\n",
        "\n",
        "def evaluate_model(interpreter):\n",
        "  input_index = interpreter.get_input_details()[0][\"index\"]\n",
        "  output_index = interpreter.get_output_details()[0][\"index\"]\n",
        "\n",
        "  # Run predictions on every image in the \"test\" dataset.\n",
        "  prediction_digits = []\n",
        "  for i, test_image in enumerate(test_images):\n",
        "    if i % 1000 == 0:\n",
        "      print('Evaluated on {n} results so far.'.format(n=i))\n",
        "    # Pre-processing: add batch dimension and convert to float32 to match with\n",
        "    # the model's input data format.\n",
        "    test_image = np.expand_dims(test_image, axis=0).astype(np.float32)\n",
        "    interpreter.set_tensor(input_index, test_image)\n",
        "\n",
        "    # Run inference.\n",
        "    interpreter.invoke()\n",
        "\n",
        "    # Post-processing: remove batch dimension and find the digit with highest\n",
        "    # probability.\n",
        "    output = interpreter.tensor(output_index)\n",
        "    digit = np.argmax(output()[0])\n",
        "    prediction_digits.append(digit)\n",
        "\n",
        "  print('\\n')\n",
        "  # Compare prediction results with ground truth labels to calculate accuracy.\n",
        "  prediction_digits = np.array(prediction_digits)\n",
        "  accuracy = (prediction_digits == test_labels).mean()\n",
        "  return accuracy"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tz43mHbZ3VHW",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 255
        },
        "outputId": "255bd574-4c59-4541-deb0-c1dc726bf9b0"
      },
      "source": [
        "# You evaluate the quantized model and see that the accuracy from TensorFlow persists to the TFLite backend.\n",
        "\n",
        "interpreter = tf.lite.Interpreter(model_content=quantized_tflite_model)\n",
        "interpreter.allocate_tensors()\n",
        "\n",
        "test_accuracy = evaluate_model(interpreter)\n",
        "\n",
        "print('Quant TFLite test_accuracy:', test_accuracy)\n",
        "print('Quant TF test accuracy:', q_aware_model_accuracy)"
      ],
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Evaluated on 0 results so far.\n",
            "Evaluated on 1000 results so far.\n",
            "Evaluated on 2000 results so far.\n",
            "Evaluated on 3000 results so far.\n",
            "Evaluated on 4000 results so far.\n",
            "Evaluated on 5000 results so far.\n",
            "Evaluated on 6000 results so far.\n",
            "Evaluated on 7000 results so far.\n",
            "Evaluated on 8000 results so far.\n",
            "Evaluated on 9000 results so far.\n",
            "\n",
            "\n",
            "Quant TFLite test_accuracy: 0.9628\n",
            "Quant TF test accuracy: 0.9682000279426575\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "h6_AQF_04sp8",
        "colab_type": "text"
      },
      "source": [
        "See 4x smaller model from quantization.   \n",
        "Compare performance of float quantization model with Quantization aware model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6hbBJPP-4f59",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 85
        },
        "outputId": "e64d1764-cf43-4d2f-9610-f982991adff5"
      },
      "source": [
        "# Create float TFLite model.\n",
        "float_converter = tf.lite.TFLiteConverter.from_keras_model(model)\n",
        "float_tflite_model = float_converter.convert()\n",
        "\n",
        "# Measure sizes of models.\n",
        "_, float_file = tempfile.mkstemp('.tflite')\n",
        "_, quant_file = tempfile.mkstemp('.tflite')\n",
        "\n",
        "with open(quant_file, 'wb') as f:\n",
        "  f.write(quantized_tflite_model)\n",
        "\n",
        "with open(float_file, 'wb') as f:\n",
        "  f.write(float_tflite_model)\n",
        "\n",
        "print(\"Float model in Mb:\", os.path.getsize(float_file) / float(2**20))\n",
        "print(\"Quantized model in Mb:\", os.path.getsize(quant_file) / float(2**20))"
      ],
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Assets written to: /tmp/tmp46gtgctw/assets\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Assets written to: /tmp/tmp46gtgctw/assets\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Float model in Mb: 0.08069229125976562\n",
            "Quantized model in Mb: 0.0236663818359375\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gWRYMHNX5FCi",
        "colab_type": "text"
      },
      "source": [
        "**Conclusions**.  \n",
        "\n",
        "\n",
        "In this tutorial, you saw how to create quantization aware models with the \n",
        "\n",
        "\n",
        "\n",
        "*   Created quantization aware training model using TF model optimization toolkit, and compared size of the model when compared to float quantization model. Quantization aware model was 4x smaller than float model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OKs1aCeH4_8d",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}