{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "(tut_three)=\n",
    "\n",
    "# **Partial Network Quantization: Specific Layer Class**\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 'Dense' layer class.\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",
    "    Specific layer classes to quantize are passed to `quantize_model()` via a `QuantizationSpec` object.\n",
    "    For layer `l`, the class name can be found using `l.__class__.__name__`. </br>\n",
    "\n",
    "```\n",
    "Refer to the [Python API](qmodel_api) documentation for more details.\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "#\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",
    "import tiny_resnet\n",
    "from tensorflow_quantization import utils\n",
    "import os\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_specific_class\")"
   ]
  },
  {
   "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_specific_class.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 [==============================] - 17s 10ms/step - loss: 1.7871 - accuracy: 0.3526 - val_loss: 1.5601 - val_accuracy: 0.4448\n",
      "Epoch 2/10\n",
      "1407/1407 [==============================] - 14s 10ms/step - loss: 1.4970 - accuracy: 0.4641 - val_loss: 1.4441 - val_accuracy: 0.4812\n",
      "Epoch 3/10\n",
      "1407/1407 [==============================] - 13s 9ms/step - loss: 1.3885 - accuracy: 0.5040 - val_loss: 1.3627 - val_accuracy: 0.5178\n",
      "Epoch 4/10\n",
      "1407/1407 [==============================] - 13s 10ms/step - loss: 1.3101 - accuracy: 0.5347 - val_loss: 1.3018 - val_accuracy: 0.5332\n",
      "Epoch 5/10\n",
      "1407/1407 [==============================] - 13s 9ms/step - loss: 1.2473 - accuracy: 0.5591 - val_loss: 1.2233 - val_accuracy: 0.5650\n",
      "Epoch 6/10\n",
      "1407/1407 [==============================] - 14s 10ms/step - loss: 1.1926 - accuracy: 0.5796 - val_loss: 1.2065 - val_accuracy: 0.5818\n",
      "Epoch 7/10\n",
      "1407/1407 [==============================] - 14s 10ms/step - loss: 1.1475 - accuracy: 0.5972 - val_loss: 1.1449 - val_accuracy: 0.5966\n",
      "Epoch 8/10\n",
      "1407/1407 [==============================] - 13s 10ms/step - loss: 1.1041 - accuracy: 0.6126 - val_loss: 1.1292 - val_accuracy: 0.6048\n",
      "Epoch 9/10\n",
      "1407/1407 [==============================] - 14s 10ms/step - loss: 1.0636 - accuracy: 0.6275 - val_loss: 1.1122 - val_accuracy: 0.6112\n",
      "Epoch 10/10\n",
      "1407/1407 [==============================] - 13s 10ms/step - loss: 1.0268 - accuracy: 0.6406 - val_loss: 1.0829 - val_accuracy: 0.6244\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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Baseline FP32 model test accuracy: 61.51\n"
     ]
    }
   ],
   "source": [
    "# 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": 6,
   "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/tutorial_onnx_models/simple_network_quantize_specific_class/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_specific_class.fp32_saved_model)\n",
    "\n",
    "# Convert FP32 model to ONNX\n",
    "utils.convert_saved_model_to_onnx(saved_model_dir = assets.simple_network_quantize_specific_class.fp32_saved_model, onnx_model_path = assets.simple_network_quantize_specific_class.fp32_onnx_model)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[I] Layer class `Conv2D` is not quantized. Partial quantization is enabled and layer class is not in user provided QuantizationSpec class object\n",
      "[I] Layer class `ReLU` is not quantized. Partial quantization is enabled and layer class is not in user provided QuantizationSpec class object\n",
      "[I] Layer class `Conv2D` is not quantized. Partial quantization is enabled and layer class is not in user provided QuantizationSpec class object\n",
      "[I] Layer class `ReLU` is not quantized. Partial quantization is enabled and layer class is not in user provided QuantizationSpec class object\n",
      "[I] Layer class `Conv2D` is not quantized. Partial quantization is enabled and layer class is not in user provided QuantizationSpec class object\n",
      "[I] Layer class `ReLU` is not quantized. Partial quantization is enabled and layer class is not in user provided QuantizationSpec class object\n",
      "[I] Layer class `Conv2D` is not quantized. Partial quantization is enabled and layer class is not in user provided QuantizationSpec class object\n",
      "[I] Layer class `Add` is not quantized. Partial quantization is enabled and layer class is not in user provided QuantizationSpec class object\n",
      "[I] Layer class `ReLU` is not quantized. Partial quantization is enabled and layer class is not in user provided QuantizationSpec class object\n",
      "[I] Layer class `Conv2D` is not quantized. Partial quantization is enabled and layer class is not in user provided QuantizationSpec class object\n",
      "[I] Layer class `ReLU` is not quantized. Partial quantization is enabled and layer class is not in user provided QuantizationSpec class object\n",
      "[I] Layer class `Conv2D` is not quantized. Partial quantization is enabled and layer class is not in user provided QuantizationSpec class object\n",
      "[I] Layer class `Conv2D` is not quantized. Partial quantization is enabled and layer class is not in user provided QuantizationSpec class object\n",
      "[I] Layer class `Add` is not quantized. Partial quantization is enabled and layer class is not in user provided QuantizationSpec class object\n",
      "[I] Layer class `ReLU` is not quantized. Partial quantization is enabled and layer class is not in user provided QuantizationSpec class object\n",
      "[I] Layer class `MaxPooling2D` is not quantized. Partial quantization is enabled and layer class is not in user provided QuantizationSpec class object\n",
      "[I] Layer class `Flatten` is not quantized. Partial quantization is enabled and layer class is not in user provided QuantizationSpec class object\n",
      "[I] Layer class `ReLU` is not quantized. Partial quantization is enabled and layer class is not in user provided QuantizationSpec class object\n"
     ]
    }
   ],
   "source": [
    "# Quantize model\n",
    "# 1.1 Create a list with keras layer classes to quantize\n",
    "qspec = QuantizationSpec()\n",
    "qspec.add(name=\"Dense\", is_keras_class=True)\n",
    "# 1.2 Call quantize model function\n",
    "q_nn_model = quantize_model(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": 8,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test accuracy immediately after quantization:60.28, diff:1.2299999999999969\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": [
    {
     "data": {
      "image/png": "[encoded data, 102372 characters]",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tf.keras.utils.plot_model(q_nn_model, to_file = assets.simple_network_quantize_specific_class.int8 + \"/model.png\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/2\n",
      "1407/1407 [==============================] - 18s 13ms/step - loss: 0.9981 - accuracy: 0.6521 - val_loss: 1.0761 - val_accuracy: 0.6324\n",
      "Epoch 2/2\n",
      "1407/1407 [==============================] - 18s 13ms/step - loss: 0.9655 - accuracy: 0.6631 - val_loss: 1.0572 - val_accuracy: 0.6302\n",
      "Accuracy after fine tuning for 2 epochs :61.82\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": 11,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:absl:Found untraced functions such as dense_layer_call_fn, dense_layer_call_and_return_conditional_losses, dense_1_layer_call_fn, dense_1_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/tutorial_onnx_models/simple_network_quantize_specific_class/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/tutorial_onnx_models/simple_network_quantize_specific_class/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_specific_class.int8_saved_model)\n",
    "\n",
    "# Convert INT8 model to ONNX\n",
    "utils.convert_saved_model_to_onnx(saved_model_dir = assets.simple_network_quantize_specific_class.int8_saved_model, onnx_model_path = assets.simple_network_quantize_specific_class.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
}
