Quantization_in_Machine_Learning.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "HgDq5C49-LQt",
        "outputId": "d5b02761-4dcf-4f55-d5cc-c6d749dcbbf7"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[?25l     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m0.0/242.5 kB\u001b[0m \u001b[31m?\u001b[0m eta \u001b[36m-:--:--\u001b[0m\r\u001b[2K     \u001b[91m━━━━━━━━━━━━━━━━\u001b[0m\u001b[91m╸\u001b[0m\u001b[90m━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m102.4/242.5 kB\u001b[0m \u001b[31m2.8 MB/s\u001b[0m eta \u001b[36m0:00:01\u001b[0m\r\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m242.5/242.5 kB\u001b[0m \u001b[31m3.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hQuantization Tutorial\n",
            "=====================\n",
            "Quantization is a technique to downsize a trained model so that you can deploy it on EDGE devices. In this tutorial, we will:\n",
            "1. Train a hand-written digits model\n",
            "2. Export it to disk and check the size of the model\n",
            "3. Use two techniques for quantization: Post-training quantization and Quantization aware training\n",
            "\n",
            "\n",
            "Downloading data from https://storage.googleapis.com/tensorflow/tf-keras-datasets/mnist.npz\n",
            "11490434/11490434 [==============================] - 1s 0us/step\n",
            "Number of training samples: 60000\n",
            "Number of test samples: 10000\n",
            "Epoch 1/5\n",
            "1875/1875 [==============================] - 25s 3ms/step - loss: 0.2748 - accuracy: 0.9211\n",
            "Epoch 2/5\n",
            "1875/1875 [==============================] - 5s 3ms/step - loss: 0.1244 - accuracy: 0.9634\n",
            "Epoch 3/5\n",
            "1875/1875 [==============================] - 7s 4ms/step - loss: 0.0863 - accuracy: 0.9746\n",
            "Epoch 4/5\n",
            "1875/1875 [==============================] - 5s 3ms/step - loss: 0.0648 - accuracy: 0.9804\n",
            "Epoch 5/5\n",
            "1875/1875 [==============================] - 7s 4ms/step - loss: 0.0514 - accuracy: 0.9840\n",
            "313/313 [==============================] - 1s 2ms/step - loss: 0.0897 - accuracy: 0.9725\n",
            "\n",
            "Post-training Quantization\n",
            "--------------------------\n",
            "Size of non-quantized model: 320004\n",
            "Size of quantized model: 84880\n",
            "\n",
            "Quantization Aware Training\n",
            "---------------------------\n",
            "Model: \"sequential\"\n",
            "_________________________________________________________________\n",
            " Layer (type)                Output Shape              Param #   \n",
            "=================================================================\n",
            " quantize_layer (QuantizeLa  (None, 28, 28)            3         \n",
            " yer)                                                            \n",
            "                                                                 \n",
            " quant_flatten (QuantizeWra  (None, 784)               1         \n",
            " pperV2)                                                         \n",
            "                                                                 \n",
            " quant_dense (QuantizeWrapp  (None, 100)               78505     \n",
            " erV2)                                                           \n",
            "                                                                 \n",
            " quant_dense_1 (QuantizeWra  (None, 10)                1015      \n",
            " pperV2)                                                         \n",
            "                                                                 \n",
            "=================================================================\n",
            "Total params: 79524 (310.64 KB)\n",
            "Trainable params: 79510 (310.59 KB)\n",
            "Non-trainable params: 14 (56.00 Byte)\n",
            "_________________________________________________________________\n",
            "1875/1875 [==============================] - 10s 5ms/step - loss: 0.0427 - accuracy: 0.9866\n",
            "313/313 [==============================] - 1s 2ms/step - loss: 0.0786 - accuracy: 0.9740\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.10/dist-packages/tensorflow/lite/python/convert.py:953: UserWarning: Statistics for quantized inputs were expected, but not specified; continuing anyway.\n",
            "  warnings.warn(\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Size of quantization aware model: 82736\n"
          ]
        }
      ],
      "source": [
        "# Install tensorflow_model_optimization\n",
        "!pip install -q tensorflow-model-optimization\n",
        "# Import necessary libraries\n",
        "import tensorflow as tf\n",
        "from tensorflow import keras\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import tensorflow_model_optimization as tfmot\n",
        "\n",
        "# Display a title for the tutorial\n",
        "print(\"Quantization Tutorial\")\n",
        "print(\"=====================\")\n",
        "print(\"Quantization is a technique to downsize a trained model so that you can deploy it on EDGE devices. In this tutorial, we will:\")\n",
        "print(\"1. Train a hand-written digits model\")\n",
        "print(\"2. Export it to disk and check the size of the model\")\n",
        "print(\"3. Use two techniques for quantization: Post-training quantization and Quantization aware training\")\n",
        "print(\"\\n\")\n",
        "\n",
        "# Load MNIST dataset\n",
        "(X_train, y_train), (X_test, y_test) = keras.datasets.mnist.load_data()\n",
        "\n",
        "# Display number of training and test samples\n",
        "print(\"Number of training samples:\", len(X_train))\n",
        "print(\"Number of test samples:\", len(X_test))\n",
        "\n",
        "# Normalize the data\n",
        "X_train = X_train / 255.0\n",
        "X_test = X_test / 255.0\n",
        "\n",
        "# Using a Flatten layer so that we don't have to call .reshape on the input dataset\n",
        "model = keras.Sequential([\n",
        "    keras.layers.Flatten(input_shape=(28, 28)),\n",
        "    keras.layers.Dense(100, activation='relu'),\n",
        "    keras.layers.Dense(10, activation='sigmoid')\n",
        "])\n",
        "\n",
        "# Compile the model\n",
        "model.compile(optimizer='adam',\n",
        "              loss='sparse_categorical_crossentropy',\n",
        "              metrics=['accuracy'])\n",
        "\n",
        "# Train the model\n",
        "model.fit(X_train, y_train, epochs=5)\n",
        "\n",
        "# Evaluate the model\n",
        "model.evaluate(X_test, y_test)\n",
        "\n",
        "# Save the model\n",
        "model.save(\"./saved_model/\")\n",
        "\n",
        "# (1) Post-training quantization\n",
        "print(\"\\nPost-training Quantization\")\n",
        "print(\"--------------------------\")\n",
        "\n",
        "# Convert to TensorFlow Lite model without quantization\n",
        "converter = tf.lite.TFLiteConverter.from_saved_model(\"./saved_model\")\n",
        "tflite_model = converter.convert()\n",
        "\n",
        "# Convert to TensorFlow Lite model with quantization\n",
        "converter.optimizations = [tf.lite.Optimize.DEFAULT]\n",
        "tflite_quant_model = converter.convert()\n",
        "\n",
        "# Print model sizes to compare\n",
        "print(\"Size of non-quantized model:\", len(tflite_model))\n",
        "print(\"Size of quantized model:\", len(tflite_quant_model))\n",
        "\n",
        "# (2) Quantization aware training\n",
        "print(\"\\nQuantization Aware Training\")\n",
        "print(\"---------------------------\")\n",
        "\n",
        "# Quantize the model\n",
        "quantize_model = tfmot.quantization.keras.quantize_model\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",
        "# Summarize the model\n",
        "q_aware_model.summary()\n",
        "\n",
        "# Train the quantization aware model\n",
        "q_aware_model.fit(X_train, y_train, epochs=1)\n",
        "\n",
        "# Evaluate the quantization aware model\n",
        "q_aware_model.evaluate(X_test, y_test)\n",
        "\n",
        "# Convert the quantization aware model\n",
        "converter = tf.lite.TFLiteConverter.from_keras_model(q_aware_model)\n",
        "converter.optimizations = [tf.lite.Optimize.DEFAULT]\n",
        "tflite_qaware_model = converter.convert()\n",
        "\n",
        "# Print size of the quantization aware model\n",
        "print(\"Size of quantization aware model:\", len(tflite_qaware_model))\n"
      ]
    }
  ]
}