tutorials/notebooks/example_keras_imagenet.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "f8194007-6ea7-4e00-8931-a37ca2d0dd20",
      "metadata": {
        "id": "f8194007-6ea7-4e00-8931-a37ca2d0dd20"
      },
      "source": [
        "# Post Training Quantization using the Model Compression Toolkit - A Quick-Start Guide"
      ]
    },
    {
      "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/example_keras_imagenet.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 shows how to quantize a pre-trained model using the Model Compression Toolkit (MCT). We will do so by giving an example of MCT's post-training quantization. As we will see, post-training quantization is a low complexity yet effective quantization method. In this example, we will 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 will cover:\n",
        "\n",
        "1. Post-Training Quantization using MCT.\n",
        "2. Loading and preprocessing ImageNet's validation dataset.\n",
        "3. Loading and preprocessing an unlabeled representative dataset from the ImageNet trainset.\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",
        "!pip install -q 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"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "0c7fed0d-cfc8-41ee-adf1-22a98110397b",
      "metadata": {
        "id": "0c7fed0d-cfc8-41ee-adf1-22a98110397b"
      },
      "source": [
        "## Dataset preparation"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "719a8ca4-6c62-4624-a1ec-662c03dde902",
      "metadata": {
        "id": "719a8ca4-6c62-4624-a1ec-662c03dde902"
      },
      "source": [
        "Assuming we've downloaded ImageNet's training dataset to a folder, let's set the folder path:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "c9893131-0a95-4472-aa42-a73bd8d50576",
      "metadata": {
        "tags": [],
        "id": "c9893131-0a95-4472-aa42-a73bd8d50576"
      },
      "outputs": [],
      "source": [
        "TRAIN_DATASET_FOLDER = '/path/to/imagenet/training/dir'"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "028112db-3143-4fcb-96ae-e639e6476c31",
      "metadata": {
        "id": "028112db-3143-4fcb-96ae-e639e6476c31"
      },
      "source": [
        "Now, let's create two functions. The first is for preprocessing the dataset and the second is for creating an unlabeled representative dataset for quantization calibration. We will use a batch size of 50:"
      ]
    },
    {
      "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": "code",
      "execution_count": null,
      "id": "0408f624-ab68-4989-95f8-f9d327882840",
      "metadata": {
        "id": "0408f624-ab68-4989-95f8-f9d327882840"
      },
      "outputs": [],
      "source": [
        "BATCH_SIZE = 50\n",
        "n_iter=10\n",
        "\n",
        "def get_representative_dataset():\n",
        "    print('loading dataset, this may take few minutes ...')\n",
        "    dataset = tf.keras.utils.image_dataset_from_directory(\n",
        "        directory=TRAIN_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",
        "    dataset = dataset.map(lambda x, y: (imagenet_preprocess_input(x, y)))\n",
        "\n",
        "    def representative_dataset():\n",
        "        for _ in range(n_iter):\n",
        "            yield dataset.take(1).get_single_element()[0].numpy()\n",
        "\n",
        "    return representative_dataset\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": [
        "Now for the main part.\n",
        "\n",
        "First, let's load a pre-trained mobilenet-v2 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": [
        "Now, we apply post-training quantization on the model. In this example, we use the default 8-bits precision and 10 calibration iterations over the representative dataset:"
      ]
    },
    {
      "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_experimental(float_model, representative_dataset_gen)"
      ]
    },
    {
      "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, let's assume we downloaded the ImageNet validation dataset to a folder with the path below:"
      ]
    },
    {
      "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 = '/path/to/imagenet/test/dir'\n",
        "def get_validation_dataset():\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"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "fd661b39-e033-4efc-a916-f97a1642cb36",
      "metadata": {
        "id": "fd661b39-e033-4efc-a916-f97a1642cb36"
      },
      "outputs": [],
      "source": [
        "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:"
      ],
      "metadata": {
        "id": "6YjIdiRRjgkL"
      },
      "id": "6YjIdiRRjgkL"
    },
    {
      "cell_type": "code",
      "source": [
	    "import tempfile\n",
        "_, tflite_file_path = tempfile.mkstemp('.tflite') # Path of exported model\n",
        "mct.exporter.keras_export_model(model=quantized_model, save_model_path=tflite_file_path,\n",
        "                                target_platform_capabilities=target_platform_cap,\n",
        "                                serialization_format=mct.exporter.KerasExportSerializationFormat.TFLITE)\n",
        "\n",
        "_, keras_file_path = tempfile.mkstemp('.h5') # Path of exported model\n",
        "mct.exporter.keras_export_model(model=quantized_model, save_model_path=keras_file_path,\n",
        "                                target_platform_capabilities=target_platform_cap,\n",
        "                                serialization_format=mct.exporter.KerasExportSerializationFormat.KERAS_H5)"
      ],
      "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
}