qkeras-mod-explained.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a class=\"anchor\" id=\"top\"></a>\n",
    "# QKeras-Mod Explained\n",
    "Author: Luca Urbinati, PhD Student @ Politecnico di Torino, luca.urbinati@polito.it. Date: 20/01/2024, v.1.0\n",
    "***\n",
    "\n",
    "### Content of this notebook\n",
    "[Chapter 1](#ch1): how to design a quantized model (with and without fused batch normalization) starting from a Keras model using a <b>modified version of QKeras</b> [1] that quantizes weights and activations to integers implementing uniform integer quantization;\n",
    "\n",
    "[Chapter 2](#ch2): <b>compare inference results</b> between the Keras model and the quantized one;\n",
    "\n",
    "[Chapter 3](#ch3): how to <b>extract quantization factors</b> (scaling factors and zero points) from each layer of the QKeras model to behave similarly to Tensorflow Lite [2][3];\n",
    "\n",
    "[Chapter 4](#ch4): how to use <b>AutoQKeras</b> to search for the best mixed-precision integer quantized model.\n",
    "\n",
    "***\n",
    "    \n",
    "### Requirements before to start\n",
    "- Read [this QKeras tutorial](https://github.com/google/qkeras/blob/master/notebook/QKerasTutorial.ipynb) to become confident with QKeras.\n",
    "- Install the conda environment [qkeras-env.yml](https://github.com/LucaUrbinati44/qkeras-mod/blob/main/qkeras-env.yml) provided in this repo and activate it (_conda activate qkeras-env_).\n",
    "- Apply the patch to QKeras' installation to have access to the modified version of QKeras (see the [README](https://github.com/LucaUrbinati44/qkeras-mod/blob/main/README.md)). \n",
    "\n",
    "### If using this code, please cite our work\n",
    "- Luca Urbinati and Mario R. Casu, \"High-Level Design of Precision-Scalable DNN Accelerators Based on Sum-Together Multiplier\", in the review process.\n",
    "\n",
    "### References\n",
    "[1] QKeras: https://github.com/google/qkeras\n",
    "\n",
    "[2] B. Jacob et al., \"Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference,\" arXiv:1712.05877 [cs, stat], Dec. 2017. Available: http://arxiv.org/abs/1712.05877\n",
    "\n",
    "[3] Mao, Lei. \"Quantization for Neural Networks\". Lei Mao’s Log Book, May 17, 2020, https://leimao.github.io/article/Neural-Networks-Quantization/\n",
    "\n",
    "[4] M. Nagel, M. Fournarakis, R. A. Amjad, Y. Bondarenko, M. van Baalen, and T. Blankevoort, “A White Paper on Neural Network Quantization.” arXiv, Jun. 15, 2021. Available: http://arxiv.org/abs/2106.08295\n",
    "\n",
    "[5] H. Wu, P. Judd, X. Zhang, M. Isaev, and P. Micikevicius, “Integer Quantization for Deep Learning Inference: Principles and Empirical Evaluation,” arXiv:2004.09602 [cs, stat], Apr. 2020, Accessed: Dec. 22, 2021. [Online]. Available: http://arxiv.org/abs/2004.09602."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "***\n",
    "<a class=\"anchor\" id=\"ch0\"></a>\n",
    "# 0) Import libraries and data\n",
    "\n",
    "Go to next: [Ch. 1](#ch1).\n",
    "\n",
    "Go to others: [Ch. 0](#ch0), [Ch. 1](#ch1), [Ch. 2](#ch2), [Ch. 3](#ch3), [Ch. 4](#ch4).\n",
    "\n",
    "Go to [Top](#top)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "import numpy as np\n",
    "import sys\n",
    "import os\n",
    "from IPython.utils import io\n",
    "import math\n",
    "from copy import deepcopy as dc\n",
    "import pandas as pd\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "import tensorflow as tf\n",
    "from tensorflow import keras\n",
    "from tensorflow.keras.utils import to_categorical\n",
    "from tensorflow.keras.models import Sequential\n",
    "from tensorflow.keras.optimizers import Adam\n",
    "from tensorflow.keras.layers import *\n",
    "from tensorflow.keras.models import Model\n",
    "\n",
    "from qkeras import *\n",
    "from qkeras.utils import *\n",
    "from qkeras.autoqkeras import *\n",
    "\n",
    "tf.keras.backend.set_floatx('float64')\n",
    "tf.keras.backend.floatx()\n",
    "\n",
    "np.set_printoptions(threshold=sys.maxsize, precision=128, suppress=True)\n",
    "\n",
    "if tf.config.list_physical_devices('GPU') == []:\n",
    "    print(\"No GPU available\")\n",
    "else:\n",
    "    print(\"GPU available\")\n",
    "    \n",
    "# Do not care about\n",
    "#---------------------------------------------------------------------------\n",
    "#RuntimeError                              Traceback (most recent call last)\n",
    "#RuntimeError: module compiled against API version 0xe but this version of numpy is 0xd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def set_seed(seed: int = 42) -> None:\n",
    "    random.seed(seed)\n",
    "    np.random.seed(seed)\n",
    "    tf.random.set_seed(seed)\n",
    "    tf.experimental.numpy.random.seed(seed)\n",
    "    # When running on the CuDNN backend, two further options must be set\n",
    "    os.environ['TF_CUDNN_DETERMINISTIC'] = '1'\n",
    "    os.environ['TF_DETERMINISTIC_OPS'] = '1'\n",
    "    # Set a fixed value for the hash seed\n",
    "    os.environ[\"PYTHONHASHSEED\"] = str(seed)\n",
    "    print(f\"Random seed set as {seed}\")\n",
    "\n",
    "seed = 42    \n",
    "set_seed(seed)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Get dummy data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorflow.keras.datasets import mnist\n",
    "\n",
    "def get_data():\n",
    "    \n",
    "    (x_train, y_train), (x_test, y_test) = mnist.load_data()\n",
    "    x_train = x_train.reshape(x_train.shape + (1,)).astype(\"float32\")\n",
    "    x_test = x_test.reshape(x_test.shape + (1,)).astype(\"float32\")\n",
    "\n",
    "    x_train /= 256.0\n",
    "    x_test /= 256.0\n",
    "\n",
    "    x_mean = np.mean(x_train, axis=0)\n",
    "\n",
    "    x_train -= x_mean\n",
    "    x_test -= x_mean\n",
    "\n",
    "    nb_classes = np.max(y_train)+1\n",
    "    y_train = to_categorical(y_train, nb_classes)\n",
    "    y_test = to_categorical(y_test, nb_classes)\n",
    "\n",
    "    return (x_train, y_train), (x_test, y_test)\n",
    "\n",
    "\n",
    "input_width = 28\n",
    "input_channels = 1\n",
    "\n",
    "(x_train, y_train), (x_test, y_test) = get_data()\n",
    "x_train, x_val, y_train, y_val = train_test_split(x_train, y_train, test_size=0.25, random_state=seed)\n",
    "\n",
    "print(x_train.shape)\n",
    "print(y_train.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "***\n",
    "<a class=\"anchor\" id=\"ch1\"></a>\n",
    "# 1) Quantized network design\n",
    "\n",
    "Go to next: [Ch. 2](#ch2).\n",
    "\n",
    "Go to others: [Ch. 0](#ch0), [Ch. 1](#ch1), [Ch. 2](#ch2), [Ch. 3](#ch3), [Ch. 4](#ch4).\n",
    "\n",
    "Go to [Top](#top).\n",
    "\n",
    "The goal of this chapter is to design a DNN that provides quantized inputs and weights to its ```QConv2D```, ```QDepthwiseConv2D```, ```QDense``` layers, as an example.\n",
    "\n",
    "It is important to know that inside the quantized kernels of QKeras there are the corresponding TensorFlow kernels (```tf.keras.backend.conv2d()```, ```tf.keras.backend.depthwise_conv2d()```, ```tf.keras.backend.dot()```) which are floating-point kernels. \n",
    "\n",
    "When running one of these quantized kernels, QKeras partially uses the technique called [\"fake quantization\"](https://github.com/google/qkeras/issues/96#issuecomment-1210877800) that is the same technique used by Tensorflow Lite [2]. This technique consists in quantizing and dequantizing inputs and weights before running the floating-point kernel. In this way, inputs, weights (and then outputs) remain floating point numbers, but can represent quantized values only. However, there is a difference: QKeras fake-quantizes only weights and biases and does not fake-quantize the inputs, i.e. they remain \"true\" floating point numbers so they can represent any number in the floating point range (you can look at the source code of ```QConv2D()``` in [qconvolutional.py#L294](https://github.com/google/qkeras/blob/eb6e0dc86c43128c6708988d9cb54d1e106685a4/qkeras/qconvolutional.py#L294) yourself). The same holds also for the outputs: they remain in floating point because computing a kernel with floating-point inputs and fake-quantized weights gives floating-point outputs.\n",
    "\n",
    "To tackle this problem, we create a new class called ```quantized_bits_featuremap()``` which is a quantizer that implements the affine quantization mapping formula [2][3]:\n",
    "\n",
    "$$x_q = \\text{clip}\\Big( \\text{round}\\big(\\frac{1}{s} x + z\\big), \\alpha_q, \\beta_q \\Big)$$\n",
    "\n",
    "We use ```quantized_bits_featuremap(alpha=\"auto\")``` in all ```QActivation``` layers that we insert in the network before and after each ```QConv2D```, ```QDepthwiseConv2D```, ```QDense``` layer, whereas we use ```quantized_bits(alpha=\"auto\")``` to quantize weights and biases because it implements only a symmetric quantized range when ```alpha=\"auto\"``` (as written in the comment [quantizers.py#L1404](https://github.com/google/qkeras/blob/c5051b51ac5d8db7b5d235419a1538258a35a8a7/qkeras/quantizers.py#L1404)). So even if we set ```symmetric=0``` and ```keep_negative=0```, it automatically forces ```symmetric=1``` ([quantizers.py#L524](https://github.com/google/qkeras/blob/b91d8815b31f05ddf9c7b6d62381df9be72a570a/qkeras/quantizers.py#L524)) and ```keep_negative=1``` ([quantizers.py#L584](https://github.com/google/qkeras/blob/b91d8815b31f05ddf9c7b6d62381df9be72a570a/qkeras/quantizers.py#L524), [quantizers.py#L603](https://github.com/google/qkeras/blob/b91d8815b31f05ddf9c7b6d62381df9be72a570a/qkeras/quantizers.py#L524)). Thefore, as an example, ```quantized_bits(4,4,0,0,alpha='auto')``` will be treated by QKeras as ```quantized_bits(4,4,1,1,alpha='auto')```. This is not a problem because we want symmetric ranges to have their zero-points null (see later). However using ```quantized_bits(alpha=\"auto\")``` for ```QActivation``` would not be good because it will not consider the zero-point. This is why we create ```quantized_bits_featuremap()```.\n",
    "\n",
    "To apply ReLU, we place the ```ReLU``` layer after the ```QActivation``` layer. We use could have declared it directly inside ```QActivation``` as ```QActivation(\"quantized_relu(bits,integer)\")```, but ```quantized_relu()``` does not quantize with the standard affine quantization mapping formula [2][3] (shown above) which is the quantization we want to implement (in fact, the argument ```alpha=\"auto\"``` is not present in ```quantized_relu()```). We also do not pass ```ReLU``` to ```QConv2D``` in the ```activation``` argument because it would be threated as ```quantized_relu()``` as well.\n",
    "\n",
    "Finally, the ```QActivation``` layer that follows the ```ReLU``` can be used to fake-quantize the features to another bitwidth precision.\n",
    "\n",
    "Regarding the number of bits for ```quantized_bits()``` and ```quantized_bits_featuremap()```, we want that ```bits``` = ```integer``` because our target is to implement integer-only arithmetic. \n",
    "\n",
    "<br><br>\n",
    "Known issues:\n",
    "1) Using ```bits``` > 31 quantizes things with ```nan```. Why? Future work\n",
    "\n",
    "2) Always explicit the value of the keyword argument ```alpha``` of  ```quantized_bits()```, that is never leave the field blank, to avoid [strange behaviors](https://github.com/google/qkeras/issues/60)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set network hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# User settings\n",
    "\n",
    "# BATCHNORM PARAMS\n",
    "fused_batchnorm = 1\n",
    "\n",
    "# POOL PARAMS\n",
    "pool_size_list = [(4, 4)]\n",
    "\n",
    "# 2DCONV PARAMS\n",
    "filters_list = [2, 3]\n",
    "kernel_size_list = [(3, 3), (3, 3)]\n",
    "strides_list = [(1, 1), (2, 2)]\n",
    "pads_list = [\"valid\", \"same\"]\n",
    "\n",
    "# DENSE PARAMS\n",
    "units_list = [10]\n",
    "\n",
    "# QUANTIZATION PARAMS\n",
    "bit_flat = 16\n",
    "\n",
    "bits_qactiv_list  = [bit_flat, bit_flat, bit_flat, bit_flat]\n",
    "bits_qweight_list = [bit_flat, bit_flat, bit_flat, 0       ] # last value is dummy\n",
    "\n",
    "#------------------------------------------"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Case without fused BatchNorm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Original Keras model that we want to convert in QKeras\n",
    "BatchNorm is not folded"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "if fused_batchnorm == 0:\n",
    "    \n",
    "    model = tf.keras.models.Sequential([\n",
    "\n",
    "        MaxPooling2D(pool_size=pool_size_list[0], name=\"pool\"),\n",
    "\n",
    "        # Example of Conv2D without Batchnorm\n",
    "        Conv2D(filters_list[0], kernel_size_list[0], strides_list[0], pads_list[0], name=\"conv2d_0\"),\n",
    "        ReLU(name=\"relu_0\"),\n",
    "\n",
    "        # Example of Conv2D with Batchnorm\n",
    "        Conv2D(filters_list[1], kernel_size_list[1], strides_list[1], pads_list[1], name=\"conv2d_1\"),\n",
    "        BatchNormalization(name=\"bn_1\"), # BatchNorm is not folded\n",
    "        ReLU(name=\"relu_1\"),\n",
    "\n",
    "        Flatten(name=\"flatten\"),\n",
    "\n",
    "        # Example of Dense without Batchnorm\n",
    "        Dense(units_list[0], name=\"dense\"),\n",
    "        Activation(\"softmax\", name=\"softmax\")\n",
    "\n",
    "    ])\n",
    "\n",
    "    model.build((None,input_width,input_width,input_channels))\n",
    "\n",
    "    model.compile(Adam(lr=0.001), loss=['categorical_crossentropy'], metrics=['accuracy'], run_eagerly=True)\n",
    "\n",
    "    model.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## QKeras model with new activation layer \"quantized_bits_featuremap\"\n",
    "quantized_bits_featuremap(bits,integer,symmetric,keep_negative,alpha,scale_axis)\n",
    "\n",
    "To run the notebook without issues, you should apply the patch to QKeras' installtion as explained in the [README](https://github.com/LucaUrbinati44/qkeras-mod/blob/main/README.md) file of this repo."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "if fused_batchnorm == 0:\n",
    "    \n",
    "    qmodel = tf.keras.models.Sequential([\n",
    "\n",
    "        MaxPooling2D(pool_size_list[0], name=\"pool\"),\n",
    "        \n",
    "        QActivation(\"quantized_bits_featuremap(%s,%s,1,1,alpha='auto',scale_axis=0)\" % \\\n",
    "                    (bits_qactiv_list[0], bits_qactiv_list[0]), name=\"act_0\"),\n",
    "\n",
    "        # Example of Conv2D without Batchnorm\n",
    "        QConv2D(filters_list[0], kernel_size_list[0], strides_list[0], pads_list[0],\n",
    "                kernel_quantizer=\"quantized_bits(%s,%s,1,1,alpha='auto')\" % \\\n",
    "                                 (bits_qweight_list[0], bits_qweight_list[0]), \n",
    "                bias_quantizer=\"quantized_bits(31,31,1,1,alpha='auto')\", name=\"conv2d_0\"),\n",
    "                #activation=\"relu\"), # This way applies quantized_relu() that we do not want\n",
    "        ReLU(name=\"relu_0\"), # Use this instead\n",
    "        QActivation(\"quantized_bits_featuremap(%s,%s,1,1,alpha='auto',scale_axis=0)\" % \\\n",
    "                    (bits_qactiv_list[1], bits_qactiv_list[1]), name=\"act_1\"),\n",
    "\n",
    "        # Example of Conv2D with fused Batchnorm\n",
    "        #QConv2DBatchnorm(filters_list[1], kernel_size_list[1], strides_list[1], pads_list[1],\n",
    "        #                 kernel_quantizer=\"quantized_bits(%s,%s,1,1,alpha='auto')\" % \\\n",
    "        #                 (bits_qweight_list[1], bits_qweight_list[1]), \n",
    "        #                 bias_quantizer=\"quantized_bits(31,31,1,1,alpha='auto')\", name=\"conv2d_1\"),\n",
    "        QConv2D(filters_list[1], kernel_size_list[1], strides_list[1], pads_list[1],\n",
    "                kernel_quantizer=\"quantized_bits(%s,%s,1,1,alpha='auto')\" % \\\n",
    "                (bits_qweight_list[1], bits_qweight_list[1]), \n",
    "                bias_quantizer=\"quantized_bits(31,31,1,1,alpha='auto')\", name=\"conv2d_1\"),\n",
    "        BatchNormalization(name=\"bn_1\"),\n",
    "        ReLU(name=\"relu_1\"),\n",
    "        QActivation(\"quantized_bits_featuremap(%s,%s,1,1,alpha='auto',scale_axis=0)\" % \\\n",
    "                    (bits_qactiv_list[2], bits_qactiv_list[2]), name=\"act_2\"),\n",
    "\n",
    "        Flatten(name=\"flatten\"),\n",
    "\n",
    "        # Example of Dense without Batchnorm (scale_axis=2 for per-layer quantization for Dense)\n",
    "        QDense(units_list[0],\n",
    "               kernel_quantizer=\"quantized_bits(%s,%s,1,1,alpha='auto',scale_axis=2)\" % \\\n",
    "                                (bits_qweight_list[2], bits_qweight_list[2]),\n",
    "               bias_quantizer=\"quantized_bits(31,31,1,1,alpha='auto')\", name=\"dense\"),\n",
    "        QActivation(\"quantized_bits_featuremap(%s,%s,1,1,alpha='auto',scale_axis=0)\" % \\\n",
    "                    (bits_qactiv_list[3], bits_qactiv_list[3]), name=\"act_3\"),\n",
    "\n",
    "        Activation(\"softmax\", name=\"softmax\")\n",
    "\n",
    "    ])\n",
    "\n",
    "    qmodel.build((None,input_width,input_width,input_channels))\n",
    "\n",
    "    qmodel.compile(Adam(lr=0.001), loss=['categorical_crossentropy'], metrics=['accuracy'], run_eagerly=True)\n",
    "\n",
    "    qmodel.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Case with fused BN"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Original Keras model that we want to convert in QKeras\n",
    "BatchNorm is folded"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "if fused_batchnorm == 1:\n",
    "    \n",
    "    model = tf.keras.models.Sequential([\n",
    "\n",
    "        MaxPooling2D(pool_size=pool_size_list[0], name=\"pool\"),\n",
    "\n",
    "        # Example of Conv2D without Batchnorm\n",
    "        Conv2D(filters_list[0], kernel_size_list[0], strides_list[0], pads_list[0], name=\"conv2d_0\"),\n",
    "        ReLU(name=\"relu_0\"),\n",
    "\n",
    "        # Example of Conv2D with Batchnorm\n",
    "        Conv2D(filters_list[1], kernel_size_list[1], strides_list[1], pads_list[1], name=\"conv2d_1\"),\n",
    "        # BatchNormalization is not needed because folded weights of qmodel \n",
    "        # will be transfered inside Conv2D (see later cells)\n",
    "        #BatchNormalization(name=\"bn_1\"),\n",
    "        ReLU(name=\"relu_1\"),\n",
    "\n",
    "        Flatten(name=\"flatten\"),\n",
    "\n",
    "        # Example of Dense without Batchnorm\n",
    "        Dense(units_list[0], name=\"dense\"),\n",
    "        Activation(\"softmax\", name=\"softmax\")\n",
    "\n",
    "    ])\n",
    "\n",
    "    model.build((None,input_width,input_width,input_channels))\n",
    "\n",
    "    model.compile(Adam(lr=0.001), loss=['categorical_crossentropy'], metrics=['accuracy'], run_eagerly=True)\n",
    "\n",
    "    model.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## QKeras model with new activation layer \"quantized_bits_featuremap\"\n",
    "quantized_bits_featuremap(bits,integer,symmetric,keep_negative,alpha,scale_axis)\n",
    "\n",
    "To run the notebook without issues, you should apply the patch to QKeras' installtion as explained in the [README](https://github.com/LucaUrbinati44/qkeras-mod/blob/main/README.md) file of this repo."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "if fused_batchnorm == 1:\n",
    "    \n",
    "    qmodel = tf.keras.models.Sequential([\n",
    "\n",
    "        MaxPooling2D(pool_size_list[0], name=\"pool\"),\n",
    "        \n",
    "        QActivation(\"quantized_bits_featuremap(%s,%s,1,1,alpha='auto',scale_axis=0)\" % \\\n",
    "                    (bits_qactiv_list[0], bits_qactiv_list[0]), name=\"act_0\"),\n",
    "\n",
    "        # Example of Conv2D without Batchnorm\n",
    "        QConv2D(filters_list[0], kernel_size_list[0], strides_list[0], pads_list[0],\n",
    "                kernel_quantizer=\"quantized_bits(%s,%s,1,1,alpha='auto')\" % \\\n",
    "                                 (bits_qweight_list[0], bits_qweight_list[0]), \n",
    "                bias_quantizer=\"quantized_bits(31,31,1,1,alpha='auto')\", name=\"conv2d_0\"),\n",
    "                #activation=\"relu\"), # This way applies quantized_relu() that we do not want\n",
    "        ReLU(name=\"relu_0\"), # Use this instead\n",
    "        QActivation(\"quantized_bits_featuremap(%s,%s,1,1,alpha='auto',scale_axis=0)\" % \\\n",
    "                    (bits_qactiv_list[1], bits_qactiv_list[1]), name=\"act_1\"),\n",
    "\n",
    "        # Example of Conv2D with fused Batchnorm\n",
    "        QConv2DBatchnorm(filters_list[1], kernel_size_list[1], strides_list[1], pads_list[1],\n",
    "                         kernel_quantizer=\"quantized_bits(%s,%s,1,1,alpha='auto')\" % \\\n",
    "                         (bits_qweight_list[1], bits_qweight_list[1]), \n",
    "                         bias_quantizer=\"quantized_bits(31,31,1,1,alpha='auto')\", name=\"conv2d_1\"),\n",
    "        #QConv2D(filters_list[1], kernel_size_list[1], strides_list[1], pads_list[1],\n",
    "        #        kernel_quantizer=\"quantized_bits(%s,%s,1,1,alpha='auto')\" % \\\n",
    "        #        (bits_qweight_list[1], bits_qweight_list[1]), \n",
    "        #        bias_quantizer=\"quantized_bits(31,31,1,1,alpha='auto')\", name=\"conv2d_1\"),\n",
    "        #BatchNormalization(name=\"bn_1\"),\n",
    "        ReLU(name=\"relu_1\"),\n",
    "        QActivation(\"quantized_bits_featuremap(%s,%s,1,1,alpha='auto',scale_axis=0)\" % \\\n",
    "                    (bits_qactiv_list[2], bits_qactiv_list[2]), name=\"act_2\"),\n",
    "\n",
    "        Flatten(name=\"flatten\"),\n",
    "\n",
    "        # Example of Dense without Batchnorm (scale_axis=2 for per-layer quantization for Dense)\n",
    "        QDense(units_list[0],\n",
    "               kernel_quantizer=\"quantized_bits(%s,%s,1,1,alpha='auto',scale_axis=2)\" % \\\n",
    "                                (bits_qweight_list[2], bits_qweight_list[2]),\n",
    "               bias_quantizer=\"quantized_bits(31,31,1,1,alpha='auto')\", name=\"dense\"),\n",
    "        QActivation(\"quantized_bits_featuremap(%s,%s,1,1,alpha='auto',scale_axis=0)\" % \\\n",
    "                    (bits_qactiv_list[3], bits_qactiv_list[3]), name=\"act_3\"),\n",
    "\n",
    "        Activation(\"softmax\", name=\"softmax\")\n",
    "\n",
    "    ])\n",
    "\n",
    "    qmodel.build((None,input_width,input_width,input_channels))\n",
    "\n",
    "    qmodel.compile(Adam(lr=0.001), loss=['categorical_crossentropy'], metrics=['accuracy'], run_eagerly=True)\n",
    "\n",
    "    qmodel.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train, save, load qmodel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# User settings\n",
    "\n",
    "train_model = 1\n",
    "epochs = 5\n",
    "\n",
    "save_model  = 1\n",
    "    \n",
    "#------------------------------------------\n",
    "\n",
    "if train_model == 1:\n",
    "    \n",
    "    qmodel.fit(x_train, y_train, validation_data=(x_val, y_val), batch_size=512, epochs=epochs, shuffle=True)    \n",
    "\n",
    "    \n",
    "output_folder = \"./training_results/\"\n",
    "\n",
    "try:\n",
    "    os.mkdir(output_folder)\n",
    "except:\n",
    "    print(\"folder exists\")\n",
    "\n",
    "save_path = \"qmodel_fusedbatchnorm_\"+str(fused_batchnorm)\n",
    "\n",
    "if save_model == 1:\n",
    "    \n",
    "    qmodel.save_weights(output_folder+save_path, overwrite=True, save_format=\"tf\")\n",
    "    os.remove(output_folder+\"checkpoint\")\n",
    "\n",
    "    print(\"Model saved correctly\")\n",
    "\n",
    "else:\n",
    "    \n",
    "    try:\n",
    "        qmodel.load_weights(output_folder+save_path, by_name=False).expect_partial()\n",
    "        print(\"Model loaded correctly\")\n",
    "    except:\n",
    "        print(\"No model to load. Train a model first\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Transfer weights from qmodel to model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "def my_evaluate(predictions, y):\n",
    "    index_pred = np.argmax(predictions)\n",
    "    index_gold = np.argmax(y)\n",
    "    if index_pred != index_gold:\n",
    "        return 0\n",
    "    else:\n",
    "        return 1\n",
    "    \n",
    "\n",
    "qlayers = []\n",
    "for qlayer in qmodel.layers:\n",
    "    if qlayer.get_weights():\n",
    "        qlayers.append(qlayer)\n",
    "\n",
    "layers = []\n",
    "for layer in model.layers:\n",
    "    if layer.get_weights():\n",
    "        layers.append(layer)\n",
    "\n",
    "        \n",
    "for qlayer, layer in zip(qlayers, layers):\n",
    "    \n",
    "    # To disable all the printings to stdout of the functions inside this statement: \n",
    "    # https://stackoverflow.com/questions/23610585/ipython-notebook-avoid-printing-within-a-function/23611571\n",
    "    with io.capture_output(stdout=True, stderr=False) as captured:\n",
    "        \n",
    "        print(qlayer.__class__.__name__)\n",
    "\n",
    "        if qlayer.get_weights() and qlayer.__class__.__name__ not in [\"BatchNormalization\"]:\n",
    "            \n",
    "            print(qlayer.name)\n",
    "            print(\"qlayer.get_weights()[0].shape:\", qlayer.get_weights()[0].shape)\n",
    "            print(\"qlayer.get_weights()[1].shape:\", qlayer.get_weights()[1].shape)\n",
    "            print(\"layer.get_weights()[0].shape:\", layer.get_weights()[0].shape)\n",
    "            print(\"layer.get_weights()[1].shape:\", layer.get_weights()[1].shape)\n",
    "            try:\n",
    "                extracted_weights = qlayer.get_folded_weights()\n",
    "                print(\"This layer IS FOLDED\")\n",
    "            except:\n",
    "                extracted_weights = qlayer.get_weights()[0:2]\n",
    "                print(\"This layer is NOT folded\")\n",
    "            \n",
    "            print(\"layer.get_weights()[0][0] MODEL BEFORE\")\n",
    "            print(layer.get_weights()[0][0])\n",
    "            print(\"layer.get_weights()[1] MODEL BEFORE\")\n",
    "            print(layer.get_weights()[1])\n",
    "            layer.set_weights(copy.deepcopy(extracted_weights))\n",
    "            print(\"layer.get_weights()[0][0] MODEL AFTER\")\n",
    "            print(layer.get_weights()[0][0])\n",
    "            print(\"layer.get_weights()[1] MODEL AFTER\")\n",
    "            print(layer.get_weights()[1])\n",
    "            \n",
    "            try:\n",
    "                print(\"qlayer.get_folded_weights()[0][0] QMODEL\")\n",
    "                print(qlayer.get_folded_weights()[0][0])\n",
    "                print(\"qlayer.get_folded_weights()[1] QMODEL\")\n",
    "                print(qlayer.get_folded_weights()[1])\n",
    "            except:\n",
    "                print(\"qlayer.get_weights()[0][0] QMODEL\")\n",
    "                print(qlayer.get_weights()[0][0])\n",
    "                print(\"qlayer.get_weights()[1] QMODEL\")\n",
    "                print(qlayer.get_weights()[1])\n",
    "            \n",
    "            result = layer.get_weights()[0:2]\n",
    "            if not np.array_equal(layer.get_weights()[0], extracted_weights[0]) or \\\n",
    "               not np.array_equal(layer.get_weights()[1], extracted_weights[1]):\n",
    "                raise Exception(\"Transfer weights failed\")\n",
    "                \n",
    "        elif qlayer.get_weights() and qlayer.__class__.__name__ in [\"BatchNormalization\"]:\n",
    "            \n",
    "            print(qlayer.name)\n",
    "            extracted_weights = qlayer.get_weights()\n",
    "            print(\"This layer is NOT folded\")\n",
    "            \n",
    "            print(\"layer.get_weights()[0] MODEL BEFORE\")\n",
    "            print(layer.get_weights()[0])\n",
    "            print(\"layer.get_weights()[1] MODEL BEFORE\")\n",
    "            print(layer.get_weights()[1])\n",
    "            layer.set_weights(copy.deepcopy(extracted_weights))\n",
    "            print(\"layer.get_weights()[0] MODEL AFTER\")\n",
    "            print(layer.get_weights()[0])\n",
    "            print(\"layer.get_weights()[1] MODEL AFTER\")\n",
    "            print(layer.get_weights()[1])\n",
    "            \n",
    "            print(\"qlayer.get_weights()[0] QMODEL\")\n",
    "            print(qlayer.get_weights()[0])\n",
    "            print(\"qlayer.get_weights()[1] QMODEL\")\n",
    "            print(qlayer.get_weights()[1])\n",
    "            \n",
    "        print(\"------------\")\n",
    "    \n",
    "print(\"Done\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "***\n",
    "<a class=\"anchor\" id=\"ch2\"></a>\n",
    "# 2) Run inference and compare model with qmodel\n",
    "\n",
    "Go to next: [Ch. 3](#ch3).\n",
    "\n",
    "Go to others: [Ch. 0](#ch0), [Ch. 1](#ch1), [Ch. 2](#ch2), [Ch. 3](#ch3), [Ch. 4](#ch4).\n",
    "\n",
    "Go to [Top](#top)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# User settings\n",
    "\n",
    "samples_to_run = 5\n",
    "\n",
    "offset = 0\n",
    "min_samples = offset\n",
    "max_samples = offset+samples_to_run\n",
    "\n",
    "print_prediction = True\n",
    "\n",
    "#------------------------------------------\n",
    "\n",
    "tot_layers = 0\n",
    "for layer in qmodel.layers:\n",
    "    if layer.get_weights():\n",
    "        tot_layers += 1\n",
    "print(\"tot_layers:\", tot_layers)\n",
    "\n",
    "iterations = 0\n",
    "test_acc_accumulator_model = 0\n",
    "test_acc_accumulator_qmodel = 0\n",
    "\n",
    "for x, y in zip(x_test[min_samples:max_samples], y_test[min_samples:max_samples]):\n",
    "        \n",
    "    x_reshaped = x.reshape(1, x.shape[0], x.shape[1], x.shape[2])\n",
    "    \n",
    "    # Predict the samples with model, i.e. the Keras model\n",
    "    pred_model = np.asarray(model.predict(x_reshaped, batch_size=1, verbose=0), dtype=np.float64)\n",
    "\n",
    "    # Predict the samples with qmodel, i.e. the modified QKeras model\n",
    "    pred_qmodel = np.asarray(qmodel.predict(x_reshaped, batch_size=1, verbose=0), dtype=np.float64)\n",
    "    \n",
    "    # Calculate test accuracy\n",
    "    test_acc_model  = my_evaluate(pred_model,  y)\n",
    "    test_acc_qmodel = my_evaluate(pred_qmodel, y)\n",
    "\n",
    "    # Check predictions\n",
    "    if print_prediction == True:\n",
    "        \n",
    "        print(\"\\npred_model:\\n\", pred_model)\n",
    "        print(\"pred_qmodel:\\n\", pred_qmodel)\n",
    "        \n",
    "        print(\"test_acc_model:   \", test_acc_model)\n",
    "        print(\"test_acc_qmodel:  \", test_acc_qmodel)\n",
    "        \n",
    "    test_acc_accumulator_model  += test_acc_model\n",
    "    test_acc_accumulator_qmodel += test_acc_qmodel\n",
    "\n",
    "    iterations = iterations + 1\n",
    "    \n",
    "    if print_prediction == True:\n",
    "        print(\"iteration %d/%d\" % (iterations, (max_samples - min_samples)))\n",
    "        print(\"-------------------------\")\n",
    "    \n",
    "print(\"-------------------------\")\n",
    "\n",
    "print(\"TOT iterations:        \", iterations)\n",
    "print(\"TOT test_acc_model:    \", test_acc_accumulator_model/iterations)\n",
    "print(\"TOT test_acc_qmodel:   \", test_acc_accumulator_qmodel/iterations)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "***\n",
    "<a class=\"anchor\" id=\"ch3\"></a>\n",
    "# 3) Extract QKeras quantization factors\n",
    "Go to next: [Ch. 4](#ch4).\n",
    "\n",
    "Go to others: [Ch. 0](#ch0), [Ch. 1](#ch1), [Ch. 2](#ch2), [Ch. 3](#ch3), [Ch. 4](#ch4).\n",
    "\n",
    "Go to [Top](#top).\n",
    "\n",
    "This is the complete quantization formula for a matrix multiplication operation (valid also for a ```Dense``` layer) taken from [[3]#Quantized-Matrix-Multiplication-Mathematics](#https://leimao.github.io/article/Neural-Networks-Quantization/#Quantized-Matrix-Multiplication-Mathematics) (an equivalent version is Eq.7 in [2]):\n",
    "$$\\begin{align} Y_{q,i,j} &= z_Y + \\frac{s_b}{s_Y} (b_{q, j} - z_b) + \\frac{s_X s_W}{s_Y} \\Bigg[ \\bigg( \\sum_{k=1}^{p} X_{q,i,k} W_{q, k,j} \\bigg) - \\bigg( z_W \\sum_{k=1}^{p} X_{q,i,k} \\bigg) - \\bigg( z_X \\sum_{k=1}^{p} W_{q, k,j} \\bigg) + p z_X z_W\\Bigg] \\end{align}$$\n",
    "\n",
    "There are some contributions that could be deleted if the <b>zero points</b> of weights ```z_w``` and biases ```z_b``` <b>are forced to be zero, i.e. if both the quantized range and the fake-quantized/floating-point range of weights and biases, respectively, are symmetric</b>. When this happens, affine quantization mapping is called scale quantization mapping [3]. It is relatively easy to set the quantized range to be symmetric (for example, in QKeras we just need to pass ```quantized_bits()``` with ```symmetric=1``` and ```keep_negative=1``` to both the arguments ```kernel_quantizer``` and ```bias_quantizer``` of each QKeras layer), but this is not the case for the fake-quantized/floating-point range. There are two ways to make the latter symmetric:\n",
    "\n",
    "1) during training, by constraining weight and bias tensors to a given symmetric range of values;\n",
    "\n",
    "2) during training, by using a different way to calculate the scaling factor ```s```. Instead of calculating it in the standard and more general way: \n",
    "$$\\begin{align} s &= \\frac{\\beta - \\alpha}{\\beta_q - \\alpha_q}\\end{align},$$\n",
    "it can be calculated as: \n",
    "$$\\begin{align} s &= \\frac{2 * max (abs (tensor), axis=[0,1])}{\\beta_q - \\alpha_q}\\end{align},$$\n",
    "where ```tensor``` is the floating-point weight/bias tensor to be quantized, ```[alpha; beta]``` is the floating-point range (where in turns ```alpha``` and ```beta``` are the minimum and maximum values of the entire tensor, so there is only one scalar ```s``` for the entire tensor), ```[alphaq; betaq]``` is the quantized range (which depends on the number of bits we want to represent the quantized data). Regarding the operations, ```abs()``` calculates the absolute value of all the elements in ```tensor``` and ```max``` with ```axis=[0,1]``` extracts the maximum value of the entire tensor/layer (so ```s``` is still a scalar). The second formula is more general because removes the constraint of searching for the minimum in the tensor and directly assumes that the floating-point range (numerator) is symmetric, even if it is not actually true, but in this way it avoids to constrain ```alpha``` and ```beta``` to be exactly equal and opposite.\n",
    "\n",
    "[This guy](https://stackoverflow.com/questions/69746834/tf-lite-model-force-symmetric-filter-weights-in-fully-connected-layers) tried to implement the first approach using [tf.keras.constraints](https://www.tensorflow.org/api_docs/python/tf/keras/constraints) without success; instead QKeras follows the second approach, as you can see in source code of ```quantized_bits()``` in [quantizers.py#L586](https://github.com/google/qkeras/blob/b91d8815b31f05ddf9c7b6d62381df9be72a570a/qkeras/quantizers.py#L586).\n",
    "\n",
    "In the light of the aforementioned explanation, in the next cells of this notebook we will extract and save to csv files only the following quantization parameters that will be needed for the inference phase:\n",
    "\n",
    "- ```wq``` and ```w_scale``` are the quantized weights and their scaling factors with zero-point weights = 0 (known after training);\n",
    "\n",
    "- ```bq``` and ```b_scale``` are the quantized biases and their scaling factors with zero-point biases = 0 (known after training);\n",
    "\n",
    "- ```subq``` is the third term in the squared brackets in the quantization formula above, i.e. the summation over the quantized weights multiplied by the zero point of the input features (known after training):\n",
    "$$\\bigg( z_X \\sum_{k=1}^{p} W_{q, k,j} \\bigg)$$\n",
    "\n",
    "- ```in_scale```, ```in_zeropoint``` are the scaling factors and the zero points of the activation layers (need calibration to get maximum absolute values of alpha and beta for each layer).\n",
    "\n",
    "To run the notebook without issues, you should edit the file [quantizers.py](https://github.com/google/qkeras/blob/b91d8815b31f05ddf9c7b6d62381df9be72a570a/qkeras/quantizers.py) to expose the following internal variables to the external world as attributes (to ease this step, just follow the instructions in the [README](https://github.com/LucaUrbinati44/qkeras-mod/blob/main/README.md) file of this repo):\n",
    "\n",
    "- ```m_i = K.cast_to_floatx(K.pow(2, self.integer))```;\n",
    "\n",
    "- ```scale1 = (K.max(abs(x), axis=axis, keepdims=True) * 2) / levels```.\n",
    "\n",
    "In particular, we need ```m_i``` and ```scale1``` to compute a scaling factor that matches the definition of scaling factor of TensorFlow Lite [2][3]. In fact, one might imagine that the scaling factor provided by the ```scale``` attribute of ```quantized_bits()``` is the same as the TensorFlow one: unfortunately it is not, as you can see from [quantizers.py#L608](https://github.com/google/qkeras/blob/b91d8815b31f05ddf9c7b6d62381df9be72a570a/qkeras/quantizers.py#L608). The correct scaling factor is ```scale = scale1 * m_i``` and has to be computed manually. The multiplication by ```m_i``` comes from the fact that QKeras' quantizers divide the tensor ```x``` by ```m_i``` before to start applying the quantization formulas as visible in [quantizers.py#L1412](https://github.com/google/qkeras/blob/c5051b51ac5d8db7b5d235419a1538258a35a8a7/qkeras/quantizers.py#L1412). Thefore, we have to rescale back the scaling factor ```scale1``` by multiplying it by ```m_i```.\n",
    "\n",
    "Finally, the following calculations show the <b>quantization of the weights</b>, which is a <b>per-channel</b>  approach, i.e. weights have a number of scaling factors and zero points equal to the number of output channels, while <b>activations are quantized in a per-layer fashion</b> (one scaling factor and one zero point for each feature map tensor). The difference is the use of ```scale_axis=0``` in ```quantized_bits_featuremap()``` for ```QActivation()```. The reason why per-channel quantization of activations is not implemented in QKeras, as well as in TensorFlow Lite, is because \"<i>per-channel quantization of activations is much harder to implement because we cannot factor the scale factor out of the summation and would, therefore, require rescaling the accumulator for each input channel</i>\" [4].\n",
    "\n",
    "For the weights of ```Dense``` layers we set ```quantized_bits_(scale_axis=2)``` to quantize them per-layer as in TensorFlow Lite."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_calibration_data(model, x_test, calibration_samples, csv_file_path_w, csv_file_path_a):\n",
    "\n",
    "    \"\"\"\n",
    "    Return two dictionaries: one for weights and biases, and one for activations\n",
    "    1) weight and bias: scale and zero-point are known after training (zero-point = 0)\n",
    "    2) activations: scale and bias vary with the input -->\n",
    "                   --> need calibration to get maximum absolute values of alpha and beta\n",
    "    \"\"\"\n",
    "    \n",
    "    \n",
    "    # Prepare the weights and biases dictionary\n",
    "    base_param_dict = {\n",
    "               \"w_scale\": 0,\n",
    "               \"b_scale\": 0,\n",
    "               \"subq\": [],\n",
    "               \"wq\": [],\n",
    "               \"bq\": []\n",
    "    }\n",
    "    \n",
    "    w_layers = []\n",
    "    for layer in model.layers:\n",
    "        if layer.__class__.__name__ in [\"QConv2D\", \"QConv2DBatchnorm\", \n",
    "                                        \"QDepthwiseConv2D\", \"QDepthwiseConv2DBatchnorm\", \n",
    "                                        \"QDense\", \"QDenseBatchnorm\"]:\n",
    "            w_layers.append(layer.name)\n",
    "\n",
    "    # Deepcopy base dict otherwise it is always the same object\n",
    "    w_dict = {k: dc(base_param_dict) for k in w_layers}\n",
    "\n",
    "    \n",
    "    \n",
    "    # Prepare the activations dictionary\n",
    "    base_act_dict = {\n",
    "        \"alpha_of_min\": [],\n",
    "        \"beta_of_max\": [],\n",
    "        \"in_scale\": [],\n",
    "        \"in_zeropoint\":[]\n",
    "    }\n",
    "    \n",
    "    a_layers = []\n",
    "    for layer in model.layers:\n",
    "        if layer.__class__.__name__ in [\"QActivation\"]:\n",
    "            a_layers.append(layer.name)\n",
    "    \n",
    "    a_dict = {k: dc(base_act_dict) for k in a_layers}   \n",
    "                       \n",
    "    \n",
    "    \n",
    "    \n",
    "    # Extract maxium absolute values of alpha and beta of activations with calibration\n",
    "    for layer in model.layers:\n",
    "    \n",
    "        alpha_of_list = []\n",
    "        beta_of_list = []\n",
    "        \n",
    "        ##### FEATURES #####\n",
    "        if layer.__class__.__name__ in [\"QActivation\"]:\n",
    "                \n",
    "            for iter, x in enumerate(x_test[0:calibration_samples]):\n",
    "\n",
    "                data = x.reshape(1, x.shape[0], x.shape[1], x.shape[2])\n",
    "\n",
    "                pred_model = np.asarray(model.predict(data, batch_size=1, verbose=0), dtype=np.float64)\n",
    "\n",
    "                quantizer_of = layer.quantizer # it is quantized_bits_featuremap\n",
    "                \n",
    "                alpha_of = quantizer_of.alpha_f.numpy().flatten()\n",
    "                beta_of = quantizer_of.beta_f.numpy().flatten()\n",
    "                alpha_of_list.append(alpha_of)\n",
    "                beta_of_list.append(beta_of)\n",
    "                        \n",
    "            alpha_of_min = np.min(alpha_of_list)\n",
    "            beta_of_max = np.max(beta_of_list)\n",
    "            \n",
    "            print(f\"layer.name: {layer.name}, \\t alpha_of_min:\", alpha_of_min, \\\n",
    "                  \"\\tbeta_of_max:\", beta_of_max)\n",
    "            \n",
    "            a_dict[layer.name][\"alpha_of_min\"].append(alpha_of_min)\n",
    "            a_dict[layer.name][\"beta_of_max\"].append(beta_of_max)\n",
    "    \n",
    "    \n",
    "    \n",
    "    # Extract weights, biases and activations\n",
    "    zeropoint_of_list = []\n",
    "    ready1 = 0\n",
    "    \n",
    "    x = x_test[0]\n",
    "    data = x.reshape(1, x.shape[0], x.shape[1], x.shape[2])\n",
    "\n",
    "    pred_model = np.asarray(model.predict(data, batch_size=1, verbose=0), dtype=np.float64)\n",
    "\n",
    "    for iter, layer in enumerate(model.layers):\n",
    "        extractor = tf.keras.Model( inputs=model.inputs,\n",
    "                                   outputs=model.get_layer(layer.name).output)\n",
    "        of = extractor(data).numpy()\n",
    "        \n",
    "        ##### WEIGHTS AND BIASES #####\n",
    "        if layer.__class__.__name__ in [\"QConv2D\", \"QConv2DBatchnorm\", \n",
    "                                        \"QDepthwiseConv2D\", \"QDepthwiseConv2DBatchnorm\", \n",
    "                                        \"QDense\", \"QDenseBatchnorm\"]:\n",
    "\n",
    "            ##### WEIGHTS #####\n",
    "            if layer.__class__.__name__ in [\"QConv2D\", \"QDepthwiseConv2D\", \"QDense\"]:\n",
    "                parameters = layer.weights\n",
    "            else:\n",
    "                parameters = layer.get_folded_weights() # folded weights not quantized\n",
    "\n",
    "            w = parameters[0].numpy()\n",
    "            quantizer_w = layer.get_quantizers()[0] # it is quantized_bits\n",
    "            alphaq_w = quantizer_w.alphaq\n",
    "            betaq_w = quantizer_w.betaq\n",
    "            scale1_w = quantizer_w.scale1.numpy().flatten()\n",
    "            m_i = quantizer_w.m_i.numpy().flatten()\n",
    "            scale_w = scale1_w * m_i # WEIGHT SCALE\n",
    "\n",
    "            if scale_w.any() != 0:\n",
    "                tmp = np.divide(w, scale_w, dtype=np.float64)\n",
    "                tmp[np.isnan(tmp)] = 0\n",
    "                tmp[np.isinf(tmp)] = 0\n",
    "                wq = np.clip(np.trunc(tmp + np.sign(tmp)*0.5), alphaq_w, betaq_w) # QUANTIZED WEIGHTS\n",
    "            else:\n",
    "                wq = np.zeros(w.size)\n",
    "                raise Exception(\"scale_w has some values equal to 0.\")\n",
    "                \n",
    "\n",
    "            # \"subq\" (part 2)\n",
    "            if layer.__class__.__name__ in [\"QConv2D\", \"QConv2DBatchnorm\"]:\n",
    "                sum_of_weights = wq.sum(axis=(0, 1, 2))\n",
    "            elif layer.__class__.__name__ in [\"QDense\", \"QDenseBatchnorm\"]:\n",
    "                sum_of_weights = wq.sum(axis=0)\n",
    "            if layer.__class__.__name__ in [\"QDepthwiseConv2D\", \"QDepthwiseConv2DBatchnorm\"]:\n",
    "                sum_of_weights = wq.sum(axis=(0, 1))\n",
    "\n",
    "            if ready1 == 1:\n",
    "                subq = (zeropoint_of_list[-1] * sum_of_weights).flatten()\n",
    "                ready1 = 0\n",
    "                w_dict[layer.name][\"subq\"] = subq\n",
    "\n",
    "\n",
    "            ##### BIASES ######\n",
    "            quantizer_b = layer.get_quantizers()[1] # it is quantized_bits\n",
    "            b = parameters[1].numpy()\n",
    "            alphaq_b = quantizer_b.alphaq\n",
    "            betaq_b = quantizer_b.betaq\n",
    "            scale1_b = quantizer_b.scale1.numpy().flatten()\n",
    "            m_i = quantizer_b.m_i.numpy().flatten()\n",
    "            scale_b = scale1_b * m_i # BIAS SCALE\n",
    "\n",
    "\n",
    "            if scale_b != 0:\n",
    "                tmp = np.divide(b, scale_b, dtype=np.float64)\n",
    "                tmp[np.isnan(tmp)] = 0\n",
    "                tmp[np.isinf(tmp)] = 0\n",
    "                bq = np.clip(np.trunc(tmp + np.sign(tmp)*0.5), alphaq_b, betaq_b) # QUANTIZED BIASES\n",
    "            else:\n",
    "                bq = np.zeros(b.size)\n",
    "\n",
    "            w_dict[layer.name]['w_scale'] = scale_w\n",
    "            w_dict[layer.name][\"b_scale\"] = scale_b\n",
    "            w_dict[layer.name][\"wq\"] = wq\n",
    "            w_dict[layer.name][\"bq\"] = bq\n",
    "\n",
    "            \n",
    "        ##### FEATURES #####\n",
    "        elif layer.__class__.__name__ in [\"QActivation\"]:\n",
    "\n",
    "            quantizer_of = layer.quantizer # it is quantized_bits_featuremap\n",
    "            alphaq_of = quantizer_of.alphaq\n",
    "            betaq_of = quantizer_of.betaq\n",
    "\n",
    "            alpha_of_min = a_dict[layer.name][\"alpha_of_min\"][0]\n",
    "            beta_of_max = a_dict[layer.name][\"beta_of_max\"][0]\n",
    "                        \n",
    "            scale_of = np.asarray([(beta_of_max - alpha_of_min) / \\\n",
    "                                   (betaq_of - alphaq_of)], dtype=np.float64)\n",
    "            scale_of[np.isnan(scale_of)] = 0\n",
    "            scale_of[np.isinf(scale_of)] = 0\n",
    "            scale_of = scale_of[0] # ACTIVATION SCALE\n",
    "\n",
    "            z_of = np.asarray([np.around(((beta_of_max*alphaq_of - alpha_of_min*betaq_of) / \\\n",
    "                                          (beta_of_max - alpha_of_min)), 0)], dtype=np.float64)\n",
    "            z_of[np.isnan(z_of)] = 0\n",
    "            z_of[np.isinf(z_of)] = 0\n",
    "            z_of = z_of[0] # ACTIVATION ZERO-POINT\n",
    "            zeropoint_of_list.append(z_of)\n",
    "            \n",
    "            ready1 = 1 # \"subq\" (part 1)\n",
    "\n",
    "            a_dict[layer.name][\"in_scale\"].append(scale_of)\n",
    "            a_dict[layer.name][\"in_zeropoint\"].append(z_of)\n",
    "\n",
    "        print(f\"{iter+1}/{len(model.layers)} \\t{layer.name}\")\n",
    "        \n",
    "        \n",
    " \n",
    "    # Index will be determined by the first layer of nested dictionaries (layers)\n",
    "    df_w = pd.DataFrame.from_dict(w_dict, orient=\"index\")\n",
    "    df_a = pd.DataFrame.from_dict(a_dict, orient=\"index\")\n",
    "    \n",
    "    df_w.to_csv(csv_file_path_w)\n",
    "    df_a.to_csv(csv_file_path_a)\n",
    "    \n",
    "    return (df_w, df_a)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# User settings\n",
    "\n",
    "calibration_samples = 10\n",
    "\n",
    "#------------------------------------------\n",
    "\n",
    "output_folder = \"./calibration_results\"\n",
    "\n",
    "try:\n",
    "    os.mkdir(output_folder)\n",
    "except:\n",
    "    print(\"folder exists\")\n",
    "    \n",
    "\n",
    "(df_w, df_a) = extract_calibration_data(qmodel, x_test, calibration_samples,\n",
    "                            output_folder+\"/extracted_weights_fused_batchnorm_\"+str(fused_batchnorm)+\".csv\",\n",
    "                            output_folder+\"/extracted_activations_fused_batchnorm_\"+str(fused_batchnorm)+\".csv\")\n",
    "\n",
    "print(\"Extraction complete\")\n",
    "\n",
    "display(df_w)\n",
    "display(df_a)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "***\n",
    "<a class=\"anchor\" id=\"ch4\"></a>\n",
    "# 4) Quantized network design for AutoQKeras\n",
    "\n",
    "Go to [Top](#top).\n",
    "\n",
    "Go to others: [Ch. 0](#ch0), [Ch. 1](#ch1), [Ch. 2](#ch2), [Ch. 3](#ch3), [Ch. 4](#ch4).\n",
    "\n",
    "To perform an hyperparameter search on a Keras model with AutoQKeras, we need to pass the Keras model to the first argument of the AutoQKeras class to create an AutoQKeras object. The input Keras model is automatically converted to a QKeras model during the building process by the ```quantize_model()``` method of AutoQKHyperModel class (see [autoqkeras_internal.py#L570](https://github.com/google/qkeras/blob/1ab354276a041b45cd72c300e89a7c51ec99fa35/qkeras/autoqkeras/autoqkeras_internal.py#L570)) (```quantize_model()``` does not accept an input model with QKeras layers).\n",
    "\n",
    "The idea is to exploit this automatic convertion to obtain a QKeras model with ```QActivation``` layers as previously described in [Ch. 1](#ch1) with as minimum changes as possible to the original Keras model definition. The changes to make to the Keras model are the following:\n",
    "\n",
    "1) Every ```Conv2D```, ```DepthwiseConv2D``` or ```Dense``` layer must be anticipated by an ```Activation``` layer with whatever activation function inside (we don't care about the type of activation because it will be replaced by AutoQKeras during the search with the ```activation``` values defined in the search space ```quantization_config```). In this example we are going to use ```sigmoid```, but any other type would be fine;\n",
    "\n",
    "2) The last ```Conv2D```, ```DepthwiseConv2D``` or ```Dense``` layer of the network must be followed by an ```Activation``` layer;\n",
    "\n",
    "3) Every ```BatchNormalization``` layer that follows a ```Conv2D``` or ```DepthwiseConv2D``` layer must be fused with the convolution. We can use the flag ```enable_bn_folding=True``` when instantiating the AutoQKeras object to automatically do the batch normalization fusion;\n",
    "\n",
    "4) Every \"real\" ```Activation``` layer (such as ```Activation(activation=\"relu\")```), or declared as argument of the ```Conv2D```, ```DepthwiseConv2D``` or ```Dense``` layer (such as ```Conv2D(activation=\"relu\")```), must be written after the ```Conv2D```, ```DepthwiseConv2D``` or ```Dense``` layer in the \"direct\" form, i.e. using a layer that has the same name of the activation (such as ```ReLU()```);\n",
    "\n",
    "5) The activations not associated to a convolutional or fully-connected layer, usually those placed as last layer of CNNs (such as softmax), must not be changed;\n",
    "\n",
    "6) Any other layer, such as ```Pooling``` or ```Add``` layers, must not be changed;\n",
    "\n",
    "7) After ```AveragePooling```, since the average of integers is not necessarily an integer, a quantization step is required (it would be better to use the same quantizer used for its input because its output range does not significantly change). Moreover, the dynamic of the activations changes between input and output of pooling. So an ```Activation``` layer is required after this layer;\n",
    "\n",
    "8) For residual connections (```Add``` layers), do as the next cell shows.\n",
    "\n",
    "To run the next cells without issues, you should edit the file [autoqkeras/autoqkeras_internal.py](https://github.com/google/qkeras/blob/master/qkeras/autoqkeras/autoqkeras_internal.py) to expose the flag ```enable_bn_folding``` to the AutoQKeras interface (externally) and to connect it to the ```AutoQKHyperModel``` class and its ```model_quantize``` method (internally). (To ease this step, just follow the instructions in the [README](https://github.com/LucaUrbinati44/qkeras-mod/blob/main/README.md) file of this repo)."
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "# Example of residual connection (\"add\" layer)\n",
    "\n",
    "...\n",
    "\n",
    "# First stack\n",
    "\n",
    "num_filters = 16\n",
    "x = Activation('sigmoid')(inputs)\n",
    "x = Conv2D(num_filters,\n",
    "            kernel_size=3,\n",
    "            strides=1,\n",
    "            padding='same',\n",
    "            kernel_initializer='he_normal',\n",
    "            kernel_regularizer=l2(1e-4))(x)\n",
    "x = BatchNormalization()(x)\n",
    "x = ReLU()(x)\n",
    "x = Activation('sigmoid')(x)\n",
    "\n",
    "y = Conv2D(num_filters,\n",
    "            kernel_size=3,\n",
    "            strides=1,\n",
    "            padding='same',\n",
    "            kernel_initializer='he_normal',\n",
    "            kernel_regularizer=l2(1e-4))(x)\n",
    "y = BatchNormalization()(y)\n",
    "y = ReLU()(y)\n",
    "y = Activation('sigmoid')(y)\n",
    "\n",
    "x = tf.keras.layers.add([x, y])\n",
    "x = ReLU()(x)\n",
    "x = Activation('sigmoid')(x)\n",
    "\n",
    "# Second stack\n",
    "\n",
    "num_filters = 32\n",
    "y = Conv2D(num_filters,\n",
    "            kernel_size=3,\n",
    "            strides=2,\n",
    "            padding='same',\n",
    "            kernel_initializer='he_normal',\n",
    "            kernel_regularizer=l2(1e-4))(x)\n",
    "...\n",
    "\n",
    "x = Conv2D(num_filters,\n",
    "        kernel_size=1,\n",
    "        strides=2,\n",
    "        padding='same',\n",
    "        kernel_initializer='he_normal',\n",
    "        kernel_regularizer=l2(1e-4))(x)\n",
    "...\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "model2 = tf.keras.models.Sequential([\n",
    "\n",
    "    MaxPooling2D(pool_size=pool_size_list[0], name=\"pool\"),\n",
    "    Activation(\"sigmoid\", name=\"act_0\"), # fake activation layer\n",
    "\n",
    "    # Example of Conv2D without Batchnorm\n",
    "    Conv2D(filters_list[0], kernel_size_list[0], strides_list[0], pads_list[0], name=\"conv2d_0\"),\n",
    "    ReLU(name=\"relu_0\"),\n",
    "    Activation(\"sigmoid\", name=\"act_1\"), # fake activation layer\n",
    "\n",
    "    # Example of Conv2D with Batchnorm\n",
    "    Conv2D(filters_list[1], kernel_size_list[1], strides_list[1], pads_list[1], name=\"conv2d_1\"),\n",
    "    # This time is needed because the folding will be carried out by AutoQKeras\n",
    "    BatchNormalization(name=\"bn_1\"),\n",
    "    ReLU(name=\"relu_1\"),\n",
    "    Activation(\"sigmoid\", name=\"act_2\"), # fake activation layer\n",
    "\n",
    "    Flatten(name=\"flatten\"),\n",
    "\n",
    "    # Example of Dense without Batchnorm\n",
    "    Dense(units_list[0], name=\"dense\"),\n",
    "    Activation(\"sigmoid\", name=\"act_3\"), # fake activation layer\n",
    "\n",
    "    Activation(\"softmax\", name=\"softmax\")\n",
    "\n",
    "])\n",
    "\n",
    "model2.build((None,input_width,input_width,input_channels))\n",
    "\n",
    "model2.compile(Adam(lr=0.001), loss=['categorical_crossentropy'], metrics=['accuracy'], run_eagerly=True)\n",
    "\n",
    "model2.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "# User settings\n",
    "\n",
    "max_trials = 3\n",
    "epochs_per_trial = 1\n",
    "\n",
    "#------------------------------------------\n",
    "\n",
    "# scale_axis=0 does a per-layer quantization, i.e. one scaling factor for the entire layer\n",
    "# otherwise qkeras automatically does a per-channel quantization, \n",
    "# i.e. one scaling facor for each channel in layer\n",
    "# we want a per-channel quantization for weigths/biases and a per-layer quantization for activations\n",
    "\n",
    "quantization_config = {\n",
    "    \"kernel\": {\n",
    "        \"quantized_bits( 4, 4,1,1,alpha='auto')\": 4,\n",
    "        \"quantized_bits( 8, 8,1,1,alpha='auto')\": 8,\n",
    "        \"quantized_bits(16,16,1,1,alpha='auto')\": 16,\n",
    "    },\n",
    "    \"bias\": {\n",
    "        \"quantized_bits(16,16,1,1,alpha='auto')\": 16,\n",
    "        \"quantized_bits(31,31,1,1,alpha='auto')\": 31,\n",
    "    },\n",
    "    \"activation\": {\n",
    "        \"quantized_bits_featuremap( 4, 4,1,1,alpha='auto',scale_axis=0)\": 4,\n",
    "        \"quantized_bits_featuremap( 8, 8,1,1,alpha='auto',scale_axis=0)\": 8,\n",
    "        \"quantized_bits_featuremap(16,16,1,1,alpha='auto',scale_axis=0)\": 16\n",
    "    }\n",
    "}\n",
    "\n",
    "# Maximum values for w, b, a\n",
    "limit = {\n",
    "    \"Dense\": [16, 31, 16],\n",
    "    \"Conv2D\": [16, 31, 16],\n",
    "    \"DepthwiseConv2D\": [16, 31, 16],\n",
    "    \"Activation\": [16],\n",
    "    \"BatchNormalization\": []\n",
    "}\n",
    "\n",
    "goal = {\n",
    "    \"type\": \"bits\",\n",
    "    \"params\": {\n",
    "        \"delta_p\": 5.0,\n",
    "        \"delta_n\": 5.0,\n",
    "        \"rate\": 2.0,\n",
    "        \"stress\": 1.0,\n",
    "        \"input_bits\": 8,\n",
    "        \"output_bits\": 8,\n",
    "        \"ref_bits\": 16,\t \n",
    "        \"config\": {\n",
    "            \"default\": [\"parameters\", \"activations\"]\n",
    "        } \n",
    "    }\n",
    "}\n",
    "\n",
    "run_config = {\n",
    "  \"output_dir\": \"./autoqkeras/autoqkeras_fused_batchnorm_\"+str(fused_batchnorm),\n",
    "  \"goal\": goal,\n",
    "  \"quantization_config\": quantization_config,\n",
    "  \"learning_rate_optimizer\": False,\n",
    "  \"transfer_weights\": True,\n",
    "  \"mode\": \"bayesian\",\n",
    "  \"seed\": seed,\n",
    "  \"limit\": limit,\n",
    "  \"tune_filters\": \"none\",\n",
    "  \"tune_filters_exceptions\": \"none\",\n",
    "  \"distribution_strategy\": tf.distribute.get_strategy(),\n",
    "  \"layer_indexes\": range(1, len(model2.layers) - 1),\n",
    "  \"max_trials\": max_trials\n",
    "}\n",
    "\n",
    "autoqk_model = AutoQKeras(model2, metrics=[\"acc\"], custom_objects={}, \n",
    "                          **run_config, overwrite=True, enable_bn_folding=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "autoqk_model.fit(x_train, y_train, validation_data=(x_val, y_val), batch_size=512, epochs=epochs_per_trial)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "qmodel2 = autoqk_model.get_best_model()  \n",
    "print_qmodel_summary(qmodel2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!--Go to next: [Ch. 5](#ch5).-->\n",
    "\n",
    "Go to others: [Ch. 0](#ch0), [Ch. 1](#ch1), [Ch. 2](#ch2), [Ch. 3](#ch3), [Ch. 4](#ch4).\n",
    "\n",
    "Go to [Top](#top)."
   ]
  }
 ],
 "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.8.13"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": false,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}