{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "(tut_two)=\n",
    "\n",
    "# **Partial Network Quantization**\n",
    "\n",
    "This example shows how the NVIDIA TensorFlow 2.x Quantization Toolkit can be used to quantize only a few layers in a\n",
    "TensorFlow 2.x model.\n",
    "\n",
    "```{eval-rst}\n",
    "\n",
    ".. admonition:: Goal\n",
    "    :class: note\n",
    "\n",
    "    #. Take a resnet-like model and train on cifar10 dataset.\n",
    "    #. Quantize only layers named 'conv2d_2' and 'dense' in the model.\n",
    "    #. Fine-tune to recover model accuracy.\n",
    "    #. Save both original and quantized model while performing ONNX conversion.\n",
    "\n",
    "```\n",
    "\n",
    "```{eval-rst}\n",
    "\n",
    ".. admonition:: Background\n",
    "    :class: tip\n",
    "\n",
    "    Few/specific layers to quantize are passed to ``quantize_model`` function as an object of ``QuantizationSpec`` class. \n",
    "    quantization mode is set to ``partial`` in ``quantize_model`` function. \n",
    "\n",
    "    Adding layers with single input to ``QuantizationSpec`` is rather simple. However, for multi-input layers, flexibility to quantize specific inputs is also provided.  \n",
    "    \n",
    "    For example, user wants to quantize layers with name ``conv2d_2`` and ``add``. \n",
    "\n",
    "    Default behavior of ``Add`` layer class is NOT to quantize any input. None of inputs to the ``Add`` class layer is quantized when following code snippet is used. \n",
    "    \n",
    "    .. code-block:: python\n",
    "\n",
    "        q_spec = QuantizationSpec()\n",
    "        layer_name = ['conv2d_2']\n",
    "        q_spec.add(name=layer_name, quantization_index=layer_quantization_index)\n",
    "\n",
    "        q_model = quantize_model(model, quantization_spec=q_spec)\n",
    "\n",
    "    \n",
    "    However, when layer of ``Add`` class is passed via ``QuantizationSpec`` object, all inputs are quantized.  \n",
    "\n",
    "    .. code-block:: python\n",
    "\n",
    "        q_spec = QuantizationSpec()\n",
    "        layer_name = ['conv2d_2', 'add']\n",
    "        q_spec.add(name=layer_name, quantization_index=layer_quantization_index)\n",
    "\n",
    "        q_model = quantize_model(model, quantization_spec=q_spec)\n",
    "\n",
    "    \n",
    "    Code to quantize input at specific index (in this case, 1) for layer ``add`` could look as follows.\n",
    "\n",
    "    .. code-block:: python\n",
    "\n",
    "        q_spec = QuantizationSpec()\n",
    "        layer_name = ['conv2d_2', 'add']\n",
    "        layer_quantization_index = [None, [1]]\n",
    "        q_spec.add(name=layer_name, quantization_index=layer_quantization_index)\n",
    "\n",
    "        q_model = quantize_model(model, quantization_spec=q_spec)\n",
    "\n",
    "    Layer name can be found from the putput of ``model.summary()`` function for Functional and Sequetial models.\n",
    "    For subclassed model, use ``KerasModelTravller`` class from tensorflow_quantization.utils.\n",
    "\n",
    "```\n",
    "Refer [Python API](qmodel_api) documentation for more details.\n",
    "\n",
    "---\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "\n",
    "#\n",
    "# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n",
    "# SPDX-License-Identifier: Apache-2.0\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",
    "#\n",
    "\n",
    "import tensorflow as tf\n",
    "from tensorflow_quantization import quantize_model, QuantizationSpec\n",
    "from tensorflow_quantization.custom_qdq_cases import ResNetV1QDQCase\n",
    "import tiny_resnet\n",
    "import os\n",
    "from tensorflow_quantization import utils\n",
    "\n",
    "tf.keras.backend.clear_session()\n",
    "\n",
    "# Create folders to save TF and ONNX models\n",
    "assets = utils.CreateAssetsFolders(os.path.join(os.getcwd(), \"tutorials\"))\n",
    "assets.add_folder(\"simple_network_quantize_partial\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Load CIFAR10 dataset\n",
    "cifar10 = tf.keras.datasets.cifar10\n",
    "(train_images, train_labels), (test_images, test_labels) = cifar10.load_data()\n",
    "\n",
    "# Normalize the input image so that each pixel value is between 0 and 1.\n",
    "train_images = train_images / 255.0\n",
    "test_images = test_images / 255.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 81200 characters]",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nn_model_original = tiny_resnet.model()\n",
    "tf.keras.utils.plot_model(nn_model_original, to_file = assets.simple_network_quantize_partial.fp32 + \"/model.png\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/10\n",
      "1407/1407 [==============================] - 18s 10ms/step - loss: 1.7617 - accuracy: 0.3615 - val_loss: 1.5624 - val_accuracy: 0.4300\n",
      "Epoch 2/10\n",
      "1407/1407 [==============================] - 14s 10ms/step - loss: 1.4876 - accuracy: 0.4645 - val_loss: 1.4242 - val_accuracy: 0.4762\n",
      "Epoch 3/10\n",
      "1407/1407 [==============================] - 14s 10ms/step - loss: 1.3737 - accuracy: 0.5092 - val_loss: 1.3406 - val_accuracy: 0.5202\n",
      "Epoch 4/10\n",
      "1407/1407 [==============================] - 14s 10ms/step - loss: 1.2952 - accuracy: 0.5396 - val_loss: 1.2768 - val_accuracy: 0.5398\n",
      "Epoch 5/10\n",
      "1407/1407 [==============================] - 14s 10ms/step - loss: 1.2370 - accuracy: 0.5649 - val_loss: 1.2466 - val_accuracy: 0.5560\n",
      "Epoch 6/10\n",
      "1407/1407 [==============================] - 14s 10ms/step - loss: 1.1857 - accuracy: 0.5812 - val_loss: 1.2052 - val_accuracy: 0.5718\n",
      "Epoch 7/10\n",
      "1407/1407 [==============================] - 14s 10ms/step - loss: 1.1442 - accuracy: 0.5972 - val_loss: 1.1836 - val_accuracy: 0.5786\n",
      "Epoch 8/10\n",
      "1407/1407 [==============================] - 14s 10ms/step - loss: 1.1079 - accuracy: 0.6091 - val_loss: 1.1356 - val_accuracy: 0.5978\n",
      "Epoch 9/10\n",
      "1407/1407 [==============================] - 14s 10ms/step - loss: 1.0702 - accuracy: 0.6220 - val_loss: 1.1244 - val_accuracy: 0.5940\n",
      "Epoch 10/10\n",
      "1407/1407 [==============================] - 14s 10ms/step - loss: 1.0354 - accuracy: 0.6368 - val_loss: 1.1019 - val_accuracy: 0.6108\n",
      "Baseline FP32 model test accuracy: 61.46\n"
     ]
    }
   ],
   "source": [
    "# Train original classification model\n",
    "nn_model_original.compile(\n",
    "    optimizer=tiny_resnet.optimizer(lr=1e-4),\n",
    "    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
    "    metrics=[\"accuracy\"],\n",
    ")\n",
    "\n",
    "nn_model_original.fit(\n",
    "    train_images, train_labels, batch_size=32, epochs=10, validation_split=0.1\n",
    ")\n",
    "\n",
    "# Get baseline model accuracy\n",
    "_, baseline_model_accuracy = nn_model_original.evaluate(\n",
    "    test_images, test_labels, verbose=0\n",
    ")\n",
    "baseline_model_accuracy = round(100 * baseline_model_accuracy, 2)\n",
    "print(\"Baseline FP32 model test accuracy:\", baseline_model_accuracy)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Assets written to: /home/sagar/nvidia/2021/Customers/Waymo/QAT/tensorrt_qat/docs/source/notebooks/tutorials/simple_network_quantize_partial/fp32/saved_model/assets\n",
      "WARNING:tensorflow:From /home/sagar/miniconda3/lib/python3.8/site-packages/tf2onnx/tf_loader.py:711: extract_sub_graph (from tensorflow.python.framework.graph_util_impl) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use `tf.compat.v1.graph_util.extract_sub_graph`\n",
      "ONNX conversion Done!\n"
     ]
    }
   ],
   "source": [
    "# save TF FP32 original model\n",
    "tf.keras.models.save_model(nn_model_original, assets.simple_network_quantize_partial.fp32_saved_model)\n",
    "\n",
    "# Convert FP32 model to ONNX\n",
    "utils.convert_saved_model_to_onnx(saved_model_dir = assets.simple_network_quantize_partial.fp32_saved_model, onnx_model_path = assets.simple_network_quantize_partial.fp32_onnx_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[I] Layer `conv2d` is not quantized. Partial quantization is enabled and layer name is not in user provided QuantizationSpec class object\n",
      "[I] Layer `re_lu` is not quantized. Partial quantization is enabled and layer name is not in user provided QuantizationSpec class object\n",
      "[I] Layer `conv2d_1` is not quantized. Partial quantization is enabled and layer name is not in user provided QuantizationSpec class object\n",
      "[I] Layer `re_lu_1` is not quantized. Partial quantization is enabled and layer name is not in user provided QuantizationSpec class object\n",
      "[I] Layer `re_lu_2` is not quantized. Partial quantization is enabled and layer name is not in user provided QuantizationSpec class object\n",
      "[I] Layer `conv2d_3` is not quantized. Partial quantization is enabled and layer name is not in user provided QuantizationSpec class object\n",
      "[I] Layer `re_lu_3` is not quantized. Partial quantization is enabled and layer name is not in user provided QuantizationSpec class object\n",
      "[I] Layer `conv2d_4` is not quantized. Partial quantization is enabled and layer name is not in user provided QuantizationSpec class object\n",
      "[I] Layer `re_lu_4` is not quantized. Partial quantization is enabled and layer name is not in user provided QuantizationSpec class object\n",
      "[I] Layer `conv2d_5` is not quantized. Partial quantization is enabled and layer name is not in user provided QuantizationSpec class object\n",
      "[I] Layer `conv2d_6` is not quantized. Partial quantization is enabled and layer name is not in user provided QuantizationSpec class object\n",
      "[I] Layer `re_lu_5` is not quantized. Partial quantization is enabled and layer name is not in user provided QuantizationSpec class object\n",
      "[I] Layer `flatten` is not quantized. Partial quantization is enabled and layer name is not in user provided QuantizationSpec class object\n",
      "[I] Layer `re_lu_6` is not quantized. Partial quantization is enabled and layer name is not in user provided QuantizationSpec class object\n",
      "[I] Layer `dense_1` is not quantized. Partial quantization is enabled and layer name is not in user provided QuantizationSpec class object\n"
     ]
    }
   ],
   "source": [
    "# Quantize model\n",
    "# 1.1 Create a dictionary to quantize only two layers named 'conv2d_2' and 'dense'\n",
    "qspec = QuantizationSpec()\n",
    "layer_name = ['conv2d_2', 'dense']\n",
    "qspec.add(name=layer_name)\n",
    "# 1.2 Call quantize model function\n",
    "q_nn_model = quantize_model(\n",
    "    model=nn_model_original, quantization_mode=\"partial\", quantization_spec=qspec)\n",
    "\n",
    "q_nn_model.compile(\n",
    "    optimizer=tiny_resnet.optimizer(lr=1e-4),\n",
    "    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
    "    metrics=[\"accuracy\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 111544 characters]",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tf.keras.utils.plot_model(q_nn_model, to_file = assets.simple_network_quantize_partial.int8 + \"/model.png\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test accuracy immediately after quantization:58.96, diff:2.5\n"
     ]
    }
   ],
   "source": [
    "_, q_model_accuracy = q_nn_model.evaluate(test_images, test_labels, verbose=0)\n",
    "q_model_accuracy = round(100 * q_model_accuracy, 2)\n",
    "print(\n",
    "    \"Test accuracy immediately after quantization:{}, diff:{}\".format(\n",
    "        q_model_accuracy, (baseline_model_accuracy - q_model_accuracy)\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/2\n",
      "1407/1407 [==============================] - 20s 14ms/step - loss: 1.0074 - accuracy: 0.6480 - val_loss: 1.0854 - val_accuracy: 0.6194\n",
      "Epoch 2/2\n",
      "1407/1407 [==============================] - 19s 14ms/step - loss: 0.9799 - accuracy: 0.6583 - val_loss: 1.0782 - val_accuracy: 0.6242\n",
      "Accuracy after fine tuning for 2 epochs :62.0\n"
     ]
    }
   ],
   "source": [
    "# Fine-tune quantized model\n",
    "fine_tune_epochs = 2\n",
    "q_nn_model.fit(\n",
    "    train_images,\n",
    "    train_labels,\n",
    "    batch_size=32,\n",
    "    epochs=fine_tune_epochs,\n",
    "    validation_split=0.1,\n",
    ")\n",
    "_, q_model_accuracy = q_nn_model.evaluate(test_images, test_labels, verbose=0)\n",
    "q_model_accuracy = round(100 * q_model_accuracy, 2)\n",
    "print(\n",
    "    \"Accuracy after fine tuning for {} epochs :{}\".format(\n",
    "        fine_tune_epochs, q_model_accuracy\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:absl:Found untraced functions such as conv2d_2_layer_call_fn, conv2d_2_layer_call_and_return_conditional_losses, dense_layer_call_fn, dense_layer_call_and_return_conditional_losses while saving (showing 4 of 4). These functions will not be directly callable after loading.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Assets written to: /home/sagar/nvidia/2021/Customers/Waymo/QAT/tensorrt_qat/docs/source/notebooks/tutorials/simple_network_quantize_partial/int8/saved_model/assets\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Assets written to: /home/sagar/nvidia/2021/Customers/Waymo/QAT/tensorrt_qat/docs/source/notebooks/tutorials/simple_network_quantize_partial/int8/saved_model/assets\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ONNX conversion Done!\n"
     ]
    }
   ],
   "source": [
    "# Save TF INT8 original model\n",
    "tf.keras.models.save_model(q_nn_model, assets.simple_network_quantize_partial.int8_saved_model)\n",
    "\n",
    "# Convert INT8 model to ONNX\n",
    "utils.convert_saved_model_to_onnx(saved_model_dir = assets.simple_network_quantize_partial.int8_saved_model, onnx_model_path = assets.simple_network_quantize_partial.int8_onnx_model)\n",
    "\n",
    "tf.keras.backend.clear_session()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "```{note}\n",
    "ONNX files can be visualized with [Netron](https://netron.app/).\n",
    "```"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "4442e1c252d743d7d1ab28567e302ebe8a15da81acb5d7e7894db75e10bdb29d"
  },
  "kernelspec": {
   "display_name": "Python 3.8.10 ('base')",
   "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.8.10"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
