tutorials/notebooks/imx500_notebooks/keras/example_keras_mobilenetv2_for_imx500.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f8194007-6ea7-4e00-8931-a37ca2d0dd20",
   "metadata": {
    "id": "f8194007-6ea7-4e00-8931-a37ca2d0dd20"
   },
   "source": [
    "# Post-Training Quantization Example of MobileNetV2 Keras Model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9be59ea8-e208-4b64-aede-1dd6270b3540",
   "metadata": {
    "id": "9be59ea8-e208-4b64-aede-1dd6270b3540"
   },
   "source": [
    "[Run this tutorial in Google Colab](https://colab.research.google.com/github/sony/model_optimization/blob/main/tutorials/notebooks/imx500_notebooks/keras/example_keras_mobilenetv2_for_imx500.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "930e6d6d-4980-4d66-beed-9ff5a494acf9",
   "metadata": {
    "id": "930e6d6d-4980-4d66-beed-9ff5a494acf9"
   },
   "source": [
    "## Overview"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "699be4fd-d382-4eec-9d3f-e2e85cfb1762",
   "metadata": {
    "id": "699be4fd-d382-4eec-9d3f-e2e85cfb1762"
   },
   "source": [
    "This tutorial demonstrates a pre-trained model quantization using the **Model Compression Toolkit (MCT)**. \n",
    "\n",
    "It is done using the MCT's **Post-Training Quantization** tool. \n",
    "\n",
    "As we will see, post-training quantization is a low complexity yet effective quantization scheme. \n",
    "\n",
    "In this example, we quantize the model and evaluate the accuracy before and after quantization."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "85199e25-c587-41b1-aaf5-e1d23ce97ca1",
   "metadata": {
    "id": "85199e25-c587-41b1-aaf5-e1d23ce97ca1"
   },
   "source": [
    "## Summary"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c0e9543-d356-412f-acf1-c2ecad553e06",
   "metadata": {
    "id": "9c0e9543-d356-412f-acf1-c2ecad553e06"
   },
   "source": [
    "In this tutorial we cover the following subjects:\n",
    "\n",
    "1. Post-Training Quantization using MCT.\n",
    "2. Loading and preprocessing ImageNet's validation dataset.\n",
    "3. Constructing an unlabeled representative dataset.\n",
    "4. Accuracy evaluation of the floating-point and the quantized models."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "04228b7c-00f1-4ded-bead-722e2a4e89a0",
   "metadata": {
    "tags": [],
    "id": "04228b7c-00f1-4ded-bead-722e2a4e89a0"
   },
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2657cf1a-654d-45a6-b877-8bf42fc26d0d",
   "metadata": {
    "id": "2657cf1a-654d-45a6-b877-8bf42fc26d0d"
   },
   "source": [
    "Install and import the relevant packages:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "324685b9-5dcc-4d22-80f4-dec9a93d3324",
   "metadata": {
    "tags": [],
    "id": "324685b9-5dcc-4d22-80f4-dec9a93d3324"
   },
   "outputs": [],
   "source": [
    "!pip install -q tensorflow\n",
    "\n",
    "import importlib\n",
    "if not importlib.util.find_spec('model_compression_toolkit'):\n",
    "    !pip install model_compression_toolkit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3f0acc8-281c-4bca-b0b9-3d7677105f19",
   "metadata": {
    "id": "b3f0acc8-281c-4bca-b0b9-3d7677105f19"
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import keras\n",
    "import model_compression_toolkit as mct\n",
    "import os"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c7fed0d-cfc8-41ee-adf1-22a98110397b",
   "metadata": {
    "id": "0c7fed0d-cfc8-41ee-adf1-22a98110397b"
   },
   "source": [
    "## Dataset preparation"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "Download ImageNet dataset with only the validation split.\n",
    "\n",
    "**Note** that for demonstration purposes we use the validation set for the model quantization routines. Usually, a subset of the training dataset is used, but loading it is a heavy procedure that is unnecessary for the sake of this demonstration.\n",
    "\n",
    "This step may take several minutes..."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "aecde59e4c37b1da"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "if not os.path.isdir('imagenet'):\n",
    "    !mkdir imagenet\n",
    "    !wget https://image-net.org/data/ILSVRC/2012/ILSVRC2012_devkit_t12.tar.gz\n",
    "    !mv ILSVRC2012_devkit_t12.tar.gz imagenet/\n",
    "    !wget https://image-net.org/data/ILSVRC/2012/ILSVRC2012_img_val.tar\n",
    "    !mv ILSVRC2012_img_val.tar imagenet/"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "5c18b26e293b085e"
  },
  {
   "cell_type": "markdown",
   "source": [
    "Extract ImageNet validation dataset using torchvision \"datasets\" module"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "e48ff22f70ce997e"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "if not os.path.isdir('imagenet/val'):\n",
    "    import torchvision\n",
    "    ds = torchvision.datasets.ImageNet(root='./imagenet', split='val')"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "5a6fb997c54aa3fb"
  },
  {
   "cell_type": "markdown",
   "id": "028112db-3143-4fcb-96ae-e639e6476c31",
   "metadata": {
    "id": "028112db-3143-4fcb-96ae-e639e6476c31"
   },
   "source": [
    "Define the required preprocessing method for the pretrained model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed56f505-97ff-4acb-8ad8-ef09c53e9d57",
   "metadata": {
    "id": "ed56f505-97ff-4acb-8ad8-ef09c53e9d57"
   },
   "outputs": [],
   "source": [
    "def imagenet_preprocess_input(images, labels):\n",
    "    return tf.keras.applications.mobilenet_v2.preprocess_input(images), labels"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Representative dataset construction\n",
    "We show how to create a generator for the representative dataset, which is required for post-training quantization.\n",
    "\n",
    "The representative dataset is used for collecting statistics on the inference outputs of all layers in the model.\n",
    " \n",
    "In order to decide on the size of the representative dataset, we configure the batch size and the number of calibration iterations.\n",
    "This gives us the total number of samples that will be used during PTQ (batch_size x n_iter).\n",
    "In this example we set `batch_size = 50` and `n_iter = 10`, resulting in a total of 500 representative images.\n",
    "\n",
    "Please ensure that the dataset path has been set correctly."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "fcbb3eecae5346a9"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0408f624-ab68-4989-95f8-f9d327882840",
   "metadata": {
    "id": "0408f624-ab68-4989-95f8-f9d327882840"
   },
   "outputs": [],
   "source": [
    "from typing import Generator\n",
    "\n",
    "REPRESENTATIVE_DATASET_FOLDER = './imagenet/val'\n",
    "BATCH_SIZE = 50\n",
    "n_iter=10\n",
    "\n",
    "# Create representative dataset generator\n",
    "def get_representative_dataset() -> Generator:\n",
    "    \"\"\"A function that loads the dataset and returns a representative dataset generator.\n",
    "\n",
    "    Returns:\n",
    "        Generator: A generator yielding batches of preprocessed images.\n",
    "    \"\"\"\n",
    "\n",
    "    # Load the dataset from folder\n",
    "    print('loading dataset, this may take few minutes ...')    \n",
    "    dataset = tf.keras.utils.image_dataset_from_directory(\n",
    "        directory=REPRESENTATIVE_DATASET_FOLDER,\n",
    "        batch_size=BATCH_SIZE,\n",
    "        image_size=[224, 224],\n",
    "        shuffle=True,\n",
    "        crop_to_aspect_ratio=True,\n",
    "        interpolation='bilinear')  \n",
    "    # Preprocess the data\n",
    "    dataset = dataset.map(lambda x, y: (imagenet_preprocess_input(x, y)))\n",
    "\n",
    "    def representative_dataset() -> Generator:\n",
    "        \"\"\"A generator function that yields batch of preprocessed images.\n",
    "\n",
    "        Yields:\n",
    "            A batch of preprocessed images.\n",
    "        \"\"\"\n",
    "        for _ in range(n_iter):\n",
    "            yield dataset.take(1).get_single_element()[0].numpy()\n",
    "\n",
    "    return representative_dataset\n",
    "\n",
    "# Create a representative dataset generator\n",
    "representative_dataset_gen = get_representative_dataset()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a1e9ba6-2954-4506-ad5c-0da273701ba5",
   "metadata": {
    "id": "4a1e9ba6-2954-4506-ad5c-0da273701ba5"
   },
   "source": [
    "## Model Post-Training quantization using MCT"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55edbb99-ab2f-4dde-aa74-4ddee61b2615",
   "metadata": {
    "id": "55edbb99-ab2f-4dde-aa74-4ddee61b2615"
   },
   "source": [
    "This is the main part in which we quantize our model.\n",
    "\n",
    "First, we load a pre-trained MobileNetV2 model from Keras, in 32-bits floating-point precision format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80cac59f-ec5e-41ca-b673-96220924a47c",
   "metadata": {
    "id": "80cac59f-ec5e-41ca-b673-96220924a47c"
   },
   "outputs": [],
   "source": [
    "from keras.applications.mobilenet_v2 import MobileNetV2\n",
    "float_model = MobileNetV2()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a8b486a-ca39-45d9-8699-f7116b0414c9",
   "metadata": {
    "id": "8a8b486a-ca39-45d9-8699-f7116b0414c9"
   },
   "source": [
    "Next, we need to define a `TargetPlatformCapability` object, representing the HW specifications on which we wish to eventually deploy our quantized model.\n",
    "\n",
    "In addition, we need to define the Quantization Configuration for our PTQ routine.\n",
    "\n",
    "Here, we demonstrate how to define a quantization configuration with several key argument that can be controlled by the user.\n",
    "**Note** that you can skip this part if you prefer to use the default quantization settings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "from model_compression_toolkit.core import QuantizationErrorMethod\n",
    "\n",
    "# Specify the IMX500-v1 target platform capability (TPC) \n",
    "tpc = mct.get_target_platform_capabilities(\"tensorflow\", 'imx500', target_platform_version='v1')\n",
    "\n",
    "# Set the following quantization configurations:\n",
    "# Choose the desired QuantizationErrorMethod for the quantization parameters search.\n",
    "# Enable weights bias correction induced by quantization.\n",
    "# Enable shift negative corrections for improving 'signed' non-linear functions quantization (such as swish, prelu, etc.) \n",
    "# Set the threshold to filter outliers with z_score of 16. \n",
    "q_config = mct.core.QuantizationConfig(activation_error_method=QuantizationErrorMethod.MSE,\n",
    "                                       weights_error_method=QuantizationErrorMethod.MSE,\n",
    "                                       weights_bias_correction=True,\n",
    "                                       shift_negative_activation_correction=True,\n",
    "                                       z_threshold=16)\n",
    "\n",
    "ptq_config = mct.core.CoreConfig(quantization_config=q_config)"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "2edacb5b7779e4d8"
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Run model Post-Training Quantization\n",
    "Lastly, we quantize our model using MCT's post-training quantization API."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "6162dd6dd1fce7ab"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33f8373a-82a5-4b97-9a10-25ee2341d148",
   "metadata": {
    "id": "33f8373a-82a5-4b97-9a10-25ee2341d148"
   },
   "outputs": [],
   "source": [
    "quantized_model, quantization_info = mct.ptq.keras_post_training_quantization(\n",
    "    in_model=float_model, \n",
    "    representative_data_gen=representative_dataset_gen, \n",
    "    core_config=ptq_config, \n",
    "    target_platform_capabilities=tpc)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7382ada6-d001-4564-907d-767fa4e9ec56",
   "metadata": {
    "id": "7382ada6-d001-4564-907d-767fa4e9ec56"
   },
   "source": [
    "That's it! Our model is now quantized."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a7a5150-3b92-49b5-abb2-06e6c5c91d6b",
   "metadata": {
    "id": "5a7a5150-3b92-49b5-abb2-06e6c5c91d6b"
   },
   "source": [
    "## Models evaluation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0ce4fc61-e13c-48be-9f7c-d441ad76a386",
   "metadata": {
    "id": "0ce4fc61-e13c-48be-9f7c-d441ad76a386"
   },
   "source": [
    "In order to evaluate our models, we first need to load the validation dataset. As before, please ensure that the dataset path has been set correctly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eef7c875-c4fc-4819-97e5-721805cba546",
   "metadata": {
    "tags": [],
    "id": "eef7c875-c4fc-4819-97e5-721805cba546"
   },
   "outputs": [],
   "source": [
    "TEST_DATASET_FOLDER = './imagenet/val'\n",
    "def get_validation_dataset() -> tf.data.Dataset:\n",
    "    \"\"\"Load the validation dataset for evaluation.\n",
    "\n",
    "    Returns:\n",
    "        tf.data.Dataset: The validation dataset.\n",
    "    \"\"\"\n",
    "    dataset = tf.keras.utils.image_dataset_from_directory(\n",
    "        directory=TEST_DATASET_FOLDER,\n",
    "        batch_size=BATCH_SIZE,\n",
    "        image_size=[224, 224],\n",
    "        shuffle=False,\n",
    "        crop_to_aspect_ratio=True,\n",
    "        interpolation='bilinear')\n",
    "    dataset = dataset.map(lambda x, y: (imagenet_preprocess_input(x, y)))\n",
    "    return dataset\n",
    "\n",
    "evaluation_dataset = get_validation_dataset()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9889d217-90a6-4615-8569-38dc9cdd5999",
   "metadata": {
    "id": "9889d217-90a6-4615-8569-38dc9cdd5999"
   },
   "source": [
    "Let's start with the floating-point model evaluation.\n",
    "\n",
    "We need to compile the model before evaluation and set the loss and the evaluation metric:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d3a0ae9-beaa-4af8-8481-49d4917c2209",
   "metadata": {
    "id": "1d3a0ae9-beaa-4af8-8481-49d4917c2209"
   },
   "outputs": [],
   "source": [
    "float_model.compile(loss=keras.losses.SparseCategoricalCrossentropy(), metrics=[\"accuracy\"])\n",
    "results = float_model.evaluate(evaluation_dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ead4a6f3-86a0-4e6c-8229-a2ff514f7b8c",
   "metadata": {
    "id": "ead4a6f3-86a0-4e6c-8229-a2ff514f7b8c"
   },
   "source": [
    "Finally, let's evaluate the quantized model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1bc377ee-39b4-4ced-95db-f7d51ab60848",
   "metadata": {
    "id": "1bc377ee-39b4-4ced-95db-f7d51ab60848"
   },
   "outputs": [],
   "source": [
    "quantized_model.compile(loss=keras.losses.SparseCategoricalCrossentropy(), metrics=[\"accuracy\"])\n",
    "results = quantized_model.evaluate(evaluation_dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ebfbb4de-5b6e-4732-83d3-a21e96cdd866",
   "metadata": {
    "id": "ebfbb4de-5b6e-4732-83d3-a21e96cdd866"
   },
   "source": [
    "You can see that we got a very small degradation with a compression rate of x4 !"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "Now, we can export the model to Keras and TFLite. Please ensure that the `save_model_path` has been set correctly."
   ],
   "metadata": {
    "id": "6YjIdiRRjgkL"
   },
   "id": "6YjIdiRRjgkL"
  },
  {
   "cell_type": "code",
   "source": [
    "mct.exporter.keras_export_model(model=quantized_model, save_model_path='./qmodel.tflite',\n",
    "                                serialization_format=mct.exporter.KerasExportSerializationFormat.TFLITE, quantization_format=mct.exporter.QuantizationFormat.FAKELY_QUANT)\n",
    "\n",
    "mct.exporter.keras_export_model(model=quantized_model, save_model_path='./qmodel.keras')"
   ],
   "metadata": {
    "id": "z3CA16-ojoFL"
   },
   "id": "z3CA16-ojoFL",
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "14877777",
   "metadata": {
    "id": "14877777"
   },
   "source": [
    "## Conclusion"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb7e1572",
   "metadata": {
    "id": "bb7e1572"
   },
   "source": [
    "In this tutorial, we demonstrated how to quantize a pre-trained model using MCT with a few lines of code. We saw that we can achieve an x4 compression ratio with minimal performance degradation.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "01c1645e-205c-4d9a-8af3-e497b3addec1",
   "metadata": {
    "id": "01c1645e-205c-4d9a-8af3-e497b3addec1"
   },
   "source": [
    "\n",
    "\n",
    "Copyright 2022 Sony Semiconductor Israel, Inc. All rights reserved.\n",
    "\n",
    "Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "you may not use this file except in compliance with the License.\n",
    "You may obtain a copy of the License at\n",
    "\n",
    "    http://www.apache.org/licenses/LICENSE-2.0\n",
    "\n",
    "Unless required by applicable law or agreed to in writing, software\n",
    "distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "See the License for the specific language governing permissions and\n",
    "limitations under the License.\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  },
  "colab": {
   "provenance": []
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}