notebook/AutoQKeras.ipynb

{
 "cells": [
   {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "##### Copyright 2020 Google LLC\n",
    "#\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",
    "# https://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."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "QC9sVuNrzT-f"
   },
   "source": [
    "# Introduction\n",
    "\n",
    "In this notebook, we show how to quantize a model using AutoQKeras.\n",
    "\n",
    "As usual, let's first make sure we are using Python 3."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "height": 51
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 926,
     "status": "ok",
     "timestamp": 1591840345558,
     "user": {
      "displayName": "Claudionor Coelho",
      "photoUrl": "",
      "userId": "01084525977535968041"
     },
     "user_tz": 420
    },
    "id": "0sY-O2IfzdB3",
    "outputId": "1c5a4e7a-1003-4b56-a30a-ca6bc196f18b"
   },
   "outputs": [],
   "source": [
    "import sys\n",
    "print(sys.version)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "6V7FxYH0zfY0"
   },
   "source": [
    "Now, let's load some packages we will need to run AutoQKeras."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "wuVqOAcbz3Go"
   },
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "import json\n",
    "import pprint\n",
    "import numpy as np\n",
    "import six\n",
    "import tempfile\n",
    "import tensorflow.compat.v2 as tf\n",
    "# V2 Behavior is necessary to use TF2 APIs before TF2 is default TF version internally.\n",
    "tf.enable_v2_behavior()\n",
    "from tensorflow.keras.optimizers import *\n",
    "\n",
    "from qkeras.autoqkeras import *\n",
    "from qkeras import *\n",
    "from qkeras.utils import model_quantize\n",
    "from qkeras.qtools import run_qtools\n",
    "from qkeras.qtools import settings as qtools_settings\n",
    "\n",
    "from tensorflow.keras.utils import to_categorical\n",
    "import tensorflow_datasets as tfds\n",
    "\n",
    "print(\"using tensorflow\", tf.__version__)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's define `get_data` and `get_model` as you may not have stand alone access to examples directory inside autoqkeras."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_data(dataset_name, fast=False):\n",
    "  \"\"\"Returns dataset from tfds.\"\"\"\n",
    "  ds_train = tfds.load(name=dataset_name, split=\"train\", batch_size=-1)\n",
    "  ds_test = tfds.load(name=dataset_name, split=\"test\", batch_size=-1)\n",
    "\n",
    "  dataset = tfds.as_numpy(ds_train)\n",
    "  x_train, y_train = dataset[\"image\"].astype(np.float32), dataset[\"label\"]\n",
    "\n",
    "  dataset = tfds.as_numpy(ds_test)\n",
    "  x_test, y_test = dataset[\"image\"].astype(np.float32), dataset[\"label\"]\n",
    "\n",
    "  if len(x_train.shape) == 3:\n",
    "    x_train = x_train.reshape(x_train.shape + (1,))\n",
    "    x_test = x_test.reshape(x_test.shape + (1,))\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",
    "  print(x_train.shape[0], \"train samples\")\n",
    "  print(x_test.shape[0], \"test samples\")\n",
    "  return (x_train, y_train), (x_test, y_test)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorflow.keras.initializers import *\n",
    "from tensorflow.keras.layers import *\n",
    "from tensorflow.keras.models import Model\n",
    "from tensorflow.keras.optimizers import *\n",
    "\n",
    "class ConvBlockNetwork(object):\n",
    "  \"\"\"Creates Convolutional block type of network.\"\"\"\n",
    "\n",
    "  def __init__(\n",
    "      self,\n",
    "      shape,\n",
    "      nb_classes,\n",
    "      kernel_size,\n",
    "      filters,\n",
    "      dropout_rate=0.0,\n",
    "      with_maxpooling=True,\n",
    "      with_batchnorm=True,\n",
    "      kernel_initializer=\"he_normal\",\n",
    "      bias_initializer=\"zeros\",\n",
    "      use_separable=False,\n",
    "      use_xnornet_trick=False,\n",
    "      all_conv=False\n",
    "  ):\n",
    "    \"\"\"Creates class.\n",
    "\n",
    "    Args:\n",
    "      shape: shape of inputs.\n",
    "      nb_classes: number of output classes.\n",
    "      kernel_size: kernel_size of network.\n",
    "      filters: sizes of filters (if entry is a list, we create a block).\n",
    "      dropout_rate: dropout rate if > 0.\n",
    "      with_maxpooling: if true, use maxpooling.\n",
    "      with_batchnorm: with BatchNormalization.\n",
    "      kernel_initializer: kernel_initializer.\n",
    "      bias_initializer: bias and beta initializer.\n",
    "      use_separable: if \"dsp\", do conv's 1x3 + 3x1. If \"mobilenet\",\n",
    "        use MobileNet separable convolution. If False or \"none\", perform single\n",
    "        conv layer.\n",
    "      use_xnornet_trick: use bn+act after max pool to enable binary\n",
    "        to avoid saturation to largest value.\n",
    "      all_conv: if true, implements all convolutional network.\n",
    "    \"\"\"\n",
    "    self.shape = shape\n",
    "    self.nb_classes = nb_classes\n",
    "    self.kernel_size = kernel_size\n",
    "    self.filters = filters\n",
    "    self.dropout_rate = dropout_rate\n",
    "    self.with_maxpooling = with_maxpooling\n",
    "    self.with_batchnorm = with_batchnorm\n",
    "    self.kernel_initializer = kernel_initializer\n",
    "    self.bias_initializer = bias_initializer\n",
    "    self.use_separable = use_separable\n",
    "    self.use_xnornet_trick = use_xnornet_trick\n",
    "    self.all_conv = all_conv\n",
    "\n",
    "  def build(self):\n",
    "    \"\"\"Builds model.\"\"\"\n",
    "    x = x_in = Input(self.shape, name=\"input\")\n",
    "    for i in range(len(self.filters)):\n",
    "      if len(self.filters) > 1:\n",
    "        name_suffix_list = [str(i)]\n",
    "      else:\n",
    "        name_suffix_list = []\n",
    "      if not isinstance(self.filters[i], list):\n",
    "        filters = [self.filters[i]]\n",
    "      else:\n",
    "        filters = self.filters[i]\n",
    "      for j in range(len(filters)):\n",
    "        if len(filters) > 1:\n",
    "          name_suffix = \"_\".join(name_suffix_list + [str(j)])\n",
    "        else:\n",
    "          name_suffix = \"_\".join(name_suffix_list)\n",
    "        if self.use_separable == \"dsp\":\n",
    "          kernels = [(1, self.kernel_size), (self.kernel_size, 1)]\n",
    "        else:\n",
    "          kernels = [(self.kernel_size, self.kernel_size)]\n",
    "        for k, kernel in enumerate(kernels):\n",
    "          strides = 1\n",
    "          if (\n",
    "              not self.with_maxpooling and j == len(filters)-1 and\n",
    "              k == len(kernels)-1\n",
    "          ):\n",
    "            strides = 2\n",
    "          if self.use_separable == \"dsp\":\n",
    "            kernel_suffix = (\n",
    "                \"\".join([str(k) for k in kernel]) + \"_\" + name_suffix)\n",
    "          elif self.use_separable == \"mobilenet\":\n",
    "            depth_suffix = (\n",
    "                \"\".join([str(k) for k in kernel]) + \"_\" + name_suffix)\n",
    "            kernel_suffix = \"11_\" + name_suffix\n",
    "          else:\n",
    "            kernel_suffix = name_suffix\n",
    "          if self.use_separable == \"mobilenet\":\n",
    "            x = DepthwiseConv2D(\n",
    "                kernel,\n",
    "                padding=\"same\", strides=strides,\n",
    "                use_bias=False,\n",
    "                name=\"conv2d_dw_\" + depth_suffix)(x)\n",
    "            if self.with_batchnorm:\n",
    "              x = BatchNormalization(name=\"conv2d_dw_bn_\" + depth_suffix)(x)\n",
    "            x = Activation(\"relu\", name=\"conv2d_dw_act_\" + depth_suffix)(x)\n",
    "            kernel = (1, 1)\n",
    "            strides = 1\n",
    "          x = Conv2D(\n",
    "              filters[j], kernel,\n",
    "              strides=strides, use_bias=not self.with_batchnorm,\n",
    "              padding=\"same\",\n",
    "              kernel_initializer=self.kernel_initializer,\n",
    "              bias_initializer=self.bias_initializer,\n",
    "              name=\"conv2d_\" + kernel_suffix)(x)\n",
    "          if not (\n",
    "              self.with_maxpooling and self.use_xnornet_trick and\n",
    "              j == len(filters)-1 and k == len(kernels)-1\n",
    "          ):\n",
    "            if self.with_batchnorm:\n",
    "              x = BatchNormalization(\n",
    "                  beta_initializer=self.bias_initializer,\n",
    "                  name=\"bn_\" + kernel_suffix)(x)\n",
    "            x = Activation(\"relu\", name=\"act_\" + kernel_suffix)(x)\n",
    "      if self.with_maxpooling:\n",
    "        x = MaxPooling2D(2, 2, name=\"mp_\" + name_suffix)(x)\n",
    "        # this is a trick from xnornet to enable full binary or ternary\n",
    "        # networks to be after maxpooling.\n",
    "        if self.use_xnornet_trick:\n",
    "          x = BatchNormalization(\n",
    "              beta_initializer=self.bias_initializer,\n",
    "              name=\"mp_bn_\" + name_suffix)(x)\n",
    "          x = Activation(\"relu\", name=\"mp_act_\" + name_suffix)(x)\n",
    "      if self.dropout_rate > 0:\n",
    "        x = Dropout(self.dropout_rate, name=\"drop_\" + name_suffix)(x)\n",
    "\n",
    "    if not self.all_conv:\n",
    "      x = Flatten(name=\"flatten\")(x)\n",
    "      x = Dense(\n",
    "          self.nb_classes,\n",
    "          kernel_initializer=self.kernel_initializer,\n",
    "          bias_initializer=self.bias_initializer,\n",
    "          name=\"dense\")(x)\n",
    "      x = Activation(\"softmax\", name=\"softmax\")(x)\n",
    "    else:\n",
    "      x = Conv2D(\n",
    "          self.nb_classes, 1, strides=1, padding=\"same\",\n",
    "          kernel_initializer=self.kernel_initializer,\n",
    "          bias_initializer=self.bias_initializer,\n",
    "          name=\"dense\")(x)\n",
    "      x = Activation(\"softmax\", name=\"softmax\")(x)\n",
    "      x = Flatten(name=\"flatten\")(x)\n",
    "\n",
    "    model = Model(inputs=[x_in], outputs=[x])\n",
    "\n",
    "    return model\n",
    "\n",
    "\n",
    "def get_model(dataset):\n",
    "  \"\"\"Returns a model for the demo of AutoQKeras.\"\"\"\n",
    "  if dataset == \"mnist\":\n",
    "    model = ConvBlockNetwork(\n",
    "        shape=(28, 28, 1),\n",
    "        nb_classes=10,\n",
    "        kernel_size=3,\n",
    "        filters=[16, 32, 48, 64, 128],\n",
    "        dropout_rate=0.2,\n",
    "        with_maxpooling=False,\n",
    "        with_batchnorm=True,\n",
    "        kernel_initializer=\"he_uniform\",\n",
    "        bias_initializer=\"zeros\",\n",
    "    ).build()\n",
    "\n",
    "  elif dataset == \"fashion_mnist\":\n",
    "    model = ConvBlockNetwork(\n",
    "        shape=(28, 28, 1),\n",
    "        nb_classes=10,\n",
    "        kernel_size=3,\n",
    "        filters=[16, [32]*3, [64]*3],\n",
    "        dropout_rate=0.2,\n",
    "        with_maxpooling=True,\n",
    "        with_batchnorm=True,\n",
    "        use_separable=\"mobilenet\",\n",
    "        kernel_initializer=\"he_uniform\",\n",
    "        bias_initializer=\"zeros\",\n",
    "        use_xnornet_trick=True\n",
    "    ).build()\n",
    "\n",
    "  elif dataset == \"cifar10\":\n",
    "    model = ConvBlockNetwork(\n",
    "        shape=(32, 32, 3),\n",
    "        nb_classes=10,\n",
    "        kernel_size=3,\n",
    "        filters=[16, [32]*3, [64]*3, [128]*3],\n",
    "        dropout_rate=0.2,\n",
    "        with_maxpooling=True,\n",
    "        with_batchnorm=True,\n",
    "        use_separable=\"mobilenet\",\n",
    "        kernel_initializer=\"he_uniform\",\n",
    "        bias_initializer=\"zeros\",\n",
    "        use_xnornet_trick=True\n",
    "    ).build()\n",
    "\n",
    "  elif dataset == \"cifar100\":\n",
    "    model = ConvBlockNetwork(\n",
    "        shape=(32, 32, 3),\n",
    "        nb_classes=100,\n",
    "        kernel_size=3,\n",
    "        filters=[16, [32]*3, [64]*3, [128]*3, [256]*3],\n",
    "        dropout_rate=0.2,\n",
    "        with_maxpooling=True,\n",
    "        with_batchnorm=True,\n",
    "        use_separable=\"mobilenet\",\n",
    "        kernel_initializer=\"he_uniform\",\n",
    "        bias_initializer=\"zeros\",\n",
    "        use_xnornet_trick=True\n",
    "    ).build()\n",
    "\n",
    "  model.summary()\n",
    "\n",
    "  return model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "uXsGtqRcN7fY"
   },
   "source": [
    "`AutoQKeras` has some examples on how to run with `mnist`, `fashion_mnist`, `cifar10` and `cifar100`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "height": 51
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 18554,
     "status": "ok",
     "timestamp": 1591840377936,
     "user": {
      "displayName": "Claudionor Coelho",
      "photoUrl": "",
      "userId": "01084525977535968041"
     },
     "user_tz": 420
    },
    "id": "lB8CBTehz9FK",
    "outputId": "09f791cf-8db5-40c5-b17d-89d433308716"
   },
   "outputs": [],
   "source": [
    "DATASET = \"mnist\"\n",
    "(x_train, y_train), (x_test, y_test) = get_data(DATASET)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "bk4rOks2OIbW"
   },
   "source": [
    "Before we create the model, let's see if we can perform distributed training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "height": 206
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 304,
     "status": "ok",
     "timestamp": 1591840378251,
     "user": {
      "displayName": "Claudionor Coelho",
      "photoUrl": "",
      "userId": "01084525977535968041"
     },
     "user_tz": 420
    },
    "id": "EMbYcKb-wMOc",
    "outputId": "22e85769-4659-4212-ccdb-4b00be2fcefe"
   },
   "outputs": [],
   "source": [
    "physical_devices = tf.config.list_physical_devices()\n",
    "for d in physical_devices:\n",
    "  print(d)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "height": 34
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 14553,
     "status": "ok",
     "timestamp": 1591840392823,
     "user": {
      "displayName": "Claudionor Coelho",
      "photoUrl": "",
      "userId": "01084525977535968041"
     },
     "user_tz": 420
    },
    "id": "eMVill0TxUuG",
    "outputId": "97c07213-fdce-4eed-9af7-cc51393cd996"
   },
   "outputs": [],
   "source": [
    "has_tpus = np.any([d.device_type == \"TPU\" for d in physical_devices])\n",
    "\n",
    "if has_tpus:\n",
    "  TPU_WORKER = 'local'\n",
    "\n",
    "  resolver = tf.distribute.cluster_resolver.TPUClusterResolver(\n",
    "      tpu=TPU_WORKER, job_name='tpu_worker')\n",
    "  if TPU_WORKER != 'local':\n",
    "    tf.config.experimental_connect_to_cluster(resolver, protocol='grpc+loas')\n",
    "  tf.tpu.experimental.initialize_tpu_system(resolver)\n",
    "  strategy = tf.distribute.experimental.TPUStrategy(resolver)\n",
    "  print('Number of devices: {}'.format(strategy.num_replicas_in_sync))\n",
    "\n",
    "  cur_strategy = strategy\n",
    "else:\n",
    "  cur_strategy = tf.distribute.get_strategy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "6FIAmXgOOPWg"
   },
   "source": [
    "Now we can create the model with the distributed strategy in place if TPUs are available. We have some test models that we can use, or you can build your own models. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "height": 977
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 1149,
     "status": "ok",
     "timestamp": 1591840393983,
     "user": {
      "displayName": "Claudionor Coelho",
      "photoUrl": "",
      "userId": "01084525977535968041"
     },
     "user_tz": 420
    },
    "id": "P0_-26kwxZiQ",
    "outputId": "bf2828fe-2968-4d7d-82e7-0e2b87f063ae"
   },
   "outputs": [],
   "source": [
    "with cur_strategy.scope():\n",
    "  model = get_model(DATASET)\n",
    "  custom_objects = {}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "jok7tJq1OVuJ"
   },
   "source": [
    "Let's see the accuracy on a unquantized model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "height": 360
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 10292,
     "status": "ok",
     "timestamp": 1591840404285,
     "user": {
      "displayName": "Claudionor Coelho",
      "photoUrl": "",
      "userId": "01084525977535968041"
     },
     "user_tz": 420
    },
    "id": "nvFSJpeDxmWZ",
    "outputId": "ceac171d-2357-4d2a-ecbe-6c2775bc2a94"
   },
   "outputs": [],
   "source": [
    "with cur_strategy.scope():\n",
    "  optimizer = Adam(lr=0.02)\n",
    "  model.compile(optimizer=optimizer, loss=\"categorical_crossentropy\", metrics=[\"acc\"])\n",
    "  model.fit(x_train, y_train, epochs=10, batch_size=2048, steps_per_epoch=29, validation_data=(x_test, y_test))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "pKArZ2VwQlph"
   },
   "source": [
    "For `mnist`, we should get 99% validation accuracy, and for `fashion_mnist`, we should get around 86% of validation accuracy. Let's get a metric for high-level estimation of energy of this model. \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "height": 1000
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 413,
     "status": "ok",
     "timestamp": 1591840404708,
     "user": {
      "displayName": "Claudionor Coelho",
      "photoUrl": "",
      "userId": "01084525977535968041"
     },
     "user_tz": 420
    },
    "id": "AlIk3gtFS6iJ",
    "outputId": "780a9c28-6234-49ff-9a85-e52bf00a5c59"
   },
   "outputs": [],
   "source": [
    "  reference_internal = \"fp32\"\n",
    "  reference_accumulator = \"fp32\"\n",
    "\n",
    "  q = run_qtools.QTools(\n",
    "      model,\n",
    "      # energy calculation using a given process\n",
    "      # \"horowitz\" refers to 45nm process published at\n",
    "      # M. Horowitz, \"1.1 Computing's energy problem (and what we can do about\n",
    "      # it), \"2014 IEEE International Solid-State Circuits Conference Digest of\n",
    "      # Technical Papers (ISSCC), San Francisco, CA, 2014, pp. 10-14, \n",
    "      # doi: 10.1109/ISSCC.2014.6757323.\n",
    "      process=\"horowitz\",\n",
    "      # quantizers for model input\n",
    "      source_quantizers=[quantized_bits(8, 0, 1)],\n",
    "      is_inference=False,\n",
    "      # absolute path (including filename) of the model weights\n",
    "      # in the future, we will attempt to optimize the power model\n",
    "      # by using weight information, although it can be used to further\n",
    "      # optimize QBatchNormalization.\n",
    "      weights_path=None,\n",
    "      # keras_quantizer to quantize weight/bias in un-quantized keras layers\n",
    "      keras_quantizer=reference_internal,\n",
    "      # keras_quantizer to quantize MAC in un-quantized keras layers\n",
    "      keras_accumulator=reference_accumulator,\n",
    "      # whether calculate baseline energy\n",
    "      for_reference=True)\n",
    "  \n",
    "# caculate energy of the derived data type map.\n",
    "energy_dict = q.pe(\n",
    "    # whether to store parameters in dram, sram, or fixed\n",
    "    weights_on_memory=\"sram\",\n",
    "    # store activations in dram or sram\n",
    "    activations_on_memory=\"sram\",\n",
    "    # minimum sram size in number of bits. Let's assume a 16MB SRAM.\n",
    "    min_sram_size=8*16*1024*1024,\n",
    "    # whether load data from dram to sram (consider sram as a cache\n",
    "    # for dram. If false, we will assume data will be already in SRAM\n",
    "    rd_wr_on_io=False)\n",
    "\n",
    "# get stats of energy distribution in each layer\n",
    "energy_profile = q.extract_energy_profile(\n",
    "    qtools_settings.cfg.include_energy, energy_dict)\n",
    "# extract sum of energy of each layer according to the rule specified in\n",
    "# qtools_settings.cfg.include_energy\n",
    "total_energy = q.extract_energy_sum(\n",
    "    qtools_settings.cfg.include_energy, energy_dict)\n",
    "\n",
    "pprint.pprint(energy_profile)\n",
    "print()\n",
    "print(\"Total energy: {:.2f} uJ\".format(total_energy / 1000000.0))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "-eDXxDSUVJ2m"
   },
   "source": [
    "During the computation, we had a dictionary that outlines the energy per layer (`energy_profile`), and total energy (`total_energy`). The reader should remember that `energy_profile` may need additional filtering as implementations will fuse some\n",
    "layers. When we compute the `total_energy`, we consider an approximation that some layers will be fused to compute the final energy number. For example, a convolution layer followed by an activation layer will be fused into a single layer so that the output of the convolution layer is not used.\n",
    "\n",
    "You have to remember that our high-level model for energy has several assumptions:\n",
    "\n",
    "The energy of a layer is estimated as `energy(layer) = energy(input) + energy(parameters) + energy(MAC) + energy(output)`.\n",
    "\n",
    "1) Reading inputs, parameters and outputs consider only _compulsory_ accesses, i.e. first access to the data, which is independent of the hardware architecture. If you remember _The 3 C's of Caches_ (https://courses.cs.washington.edu/courses/cse410/99au/lectures/Lecture-10-18/tsld035.htm) other types of accesses will depend on the accelerator architecture.\n",
    "\n",
    "2) For the multiply-and-add (MAC) energy estimation, we only consider the energy to compute the MAC, but not any other type energy. For example, in a real accelerator, you have registers, glue logic, pipeline logic that will affect the overall energy profile of the device.\n",
    "\n",
    "Although this model is simple and provides an initial estimate on what to expect, it has high-variance with respect to actual energy numbers you will find in practice, especially with respect to different architectural implementations.\n",
    "\n",
    "We assume that the real energy `Energy(layer)` is a linear combination of the high-level energy model, i.e.`Energy(layer) = k1 * energy(layer) + k2`, where `k1` and `k2` are constants that depend on the architecture of the accelerator. One can think of `k1` as the factor that accounts for the additional storage to keep the model running, and `k2` as the additional always on logic that is required to perform the operations. If we compare the energy of two implementations with different quantizations of the same layer, let's say `layer1` and `layer2`, `Energy(layer1) > Energy(layer2)` holds true iff `energy(layer1) > energy(layer2)` for the same architecture, but for different architectures, this will not be true in general.\n",
    "\n",
    "Despite its limitations to predict a single energy number, this model is quite good to compare the energy of two different models, or different types of quantizations, when we restrict it to a single architecture, and that's how we use it here."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Hr1FL8wVSy-q"
   },
   "source": [
    "# Quantizing a Model With `AutoQKeras`\n",
    "\n",
    "To quantize this model with `AutoQKeras`, we need to define the quantization for kernels, biases and activations; forgiving factors and quantization strategy.\n",
    "\n",
    "Below we define which quantizers are allowed for kernel, bias, activations and linear. Linear is a proxy that we use to capture `Activation(\"linear\")` to apply quantization without applying a non-linear operation.  In some networks, we found that this trick may be necessary to better represent the quantization space.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "vSsEwDr_yRG4"
   },
   "outputs": [],
   "source": [
    "quantization_config = {\n",
    "        \"kernel\": {\n",
    "                \"binary\": 1,\n",
    "                \"stochastic_binary\": 1,\n",
    "                \"ternary\": 2,\n",
    "                \"stochastic_ternary\": 2,\n",
    "                \"quantized_bits(2,1,1,alpha=1.0)\": 2,\n",
    "                \"quantized_bits(4,0,1,alpha=1.0)\": 4,\n",
    "                \"quantized_bits(8,0,1,alpha=1.0)\": 8,\n",
    "                \"quantized_po2(4,1)\": 4\n",
    "        },\n",
    "        \"bias\": {\n",
    "                \"quantized_bits(4,0,1)\": 4,\n",
    "                \"quantized_bits(8,3,1)\": 8,\n",
    "                \"quantized_po2(4,8)\": 4\n",
    "        },\n",
    "        \"activation\": {\n",
    "                \"binary\": 1,\n",
    "                \"ternary\": 2,\n",
    "                \"quantized_relu_po2(4,4)\": 4,\n",
    "                \"quantized_relu(3,1)\": 3,\n",
    "                \"quantized_relu(4,2)\": 4,\n",
    "                \"quantized_relu(8,2)\": 8,\n",
    "                \"quantized_relu(8,4)\": 8,\n",
    "                \"quantized_relu(16,8)\": 16\n",
    "        },\n",
    "        \"linear\": {\n",
    "                \"binary\": 1,\n",
    "                \"ternary\": 2,\n",
    "                \"quantized_bits(4,1)\": 4,\n",
    "                \"quantized_bits(8,2)\": 8,\n",
    "                \"quantized_bits(16,10)\": 16\n",
    "        }\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "GmW_xaAvZo4D"
   },
   "source": [
    "Now let's define how to apply quantization. In the simplest form, we specify how many bits for kernels, biases and activations by layer types. Note that the entry `BatchNormalization` needs to be specified here, as we only quantize layer types specified by these patterns.  For example, a `Flatten` layer is not quantized as it does not change the data type of its inputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "emTRLIZmR-P7"
   },
   "outputs": [],
   "source": [
    "limit = {\n",
    "    \"Dense\": [8, 8, 4],\n",
    "    \"Conv2D\": [4, 8, 4],\n",
    "    \"DepthwiseConv2D\": [4, 8, 4],\n",
    "    \"Activation\": [4],\n",
    "    \"BatchNormalization\": []\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "-iu5gFNhaLNE"
   },
   "source": [
    "Here, we are specifying that we want to use at most 4 bits for weights and activations, and at most 8 bits for biases in convolutional and depthwise convolutions, but we allow up to 8 bits for kernels in dense layers."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ZUMQGEIDblSa"
   },
   "source": [
    "Let's define now the forgiving factor. We will consider energy minimization as a goal as follows.  Here, we are saying that we allow 8% reduction in accuracy for a 2x reduction in energy, both reference and trials have parameters and activations on SRAM, both reference model and quantization trials do not read/write from DRAM on I/O operations, and we should consider both experiments to use SRAMs with minimum tensor sizes (commonly called distributed SRAM implementation).\n",
    "\n",
    "We also need to specify the quantizers for the inputs. In this case, we want to use `int8` as source quantizers. Other possible types are `int16`, `int32`, `fp16` or `fp32`, besides `QKeras` quantizer types.\n",
    "\n",
    "Finally, to be fair, we want to compare our quantization against fixed-point 8-bit inputs, outputs, activations, weights and biases, and 32-bit accumulators.\n",
    "\n",
    "Remember that a `forgiving factor` forgives a drop in a metric such as `accuracy` if the gains of the model are much bigger than the drop. For example, it corresponds to the sentence *we allow $\\tt{delta}\\%$ reduction in accuracy if the quantized model has $\\tt{rate} \\times$ smaller energy than the original model*, being a multiplicative factor to the metric. It is computed by $1 + \\tt{delta} \\times  \\log_{\\tt{rate}}(\\tt{stress} \\times \\tt{reference\\_cost} / \\tt{trial\\_cost})$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "kS31TuZ-aKb1"
   },
   "outputs": [],
   "source": [
    "goal = {\n",
    "    \"type\": \"energy\",\n",
    "    \"params\": {\n",
    "        \"delta_p\": 8.0,\n",
    "        \"delta_n\": 8.0,\n",
    "        \"rate\": 2.0,\n",
    "        \"stress\": 1.0,\n",
    "        \"process\": \"horowitz\",\n",
    "        \"parameters_on_memory\": [\"sram\", \"sram\"],\n",
    "        \"activations_on_memory\": [\"sram\", \"sram\"],\n",
    "        \"rd_wr_on_io\": [False, False],\n",
    "        \"min_sram_size\": [0, 0],\n",
    "        \"source_quantizers\": [\"int8\"],\n",
    "        \"reference_internal\": \"int8\",\n",
    "        \"reference_accumulator\": \"int32\"\n",
    "        }\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "-QzyWPA-dCxm"
   },
   "source": [
    "There are a few more things we need to define. Let's bundle them on a dictionary and pass them to `AutoQKeras`.  We will try a maximum of 10 trials (`max_trials`) just to limit the time we will spend finding the best quantization here.  Please note that this parameter is not valid if you are running in `hyperband` mode.\n",
    "\n",
    "`output_dir` is the directory where we will store our results. Since we are running on a colab, we will let `tempfile` chooce a directory for us.\n",
    "\n",
    "`learning_rate_optimizer` allows `AutoQKeras` to change the optimization function and the `learning_rate` to try to improve the quantization results. Since it is still experimental, it may be the case that in some cases it will get worse results. \n",
    "\n",
    "Because we are tuning filters as well, we should set `transfer_weights` to `False` as the trainable parameters will have different shapes.\n",
    "\n",
    "In `AutoQKeras` we have three modes of operation: `random`, `bayesian` and `hyperband`. I recommend the user to refer to `KerasTuner` (https://keras-team.github.io/keras-tuner/) for a complete description of them.\n",
    "\n",
    "`tune_filters` can be set to `layer`, `block` or `none`. If `tune_filters` is `block`, we change the filters by the same amount for all layers being quantized in the trial. If `tune_filters` is `layer`, we will possibly change the number of filters for each layer independently. Finally, if `tune_filters` is `none`, we will not perform filter tuning.\n",
    "\n",
    "Together with `tune_filters`, `tune_filter_exceptions` allows the user to specify by a regular expression which filters we should not perform filter tuning, which is especially good for the last layers of the network.\n",
    "\n",
    "Filter tuning is a very important feature of `AutoQKeras`. When we deep quantize a model, we may need less or more filters for each layer (and you can guess we do not know a priori how many filters we will need for each layer). Let me give you a rationale behind this.\n",
    "\n",
    "- **less filters**: let us assume we have two set of filter coefficients we want quantize: $[-0.3, 0.2, 0.5, 0.15]$ and $[-0.5, 0.4, 0.1, 0.65]$. If we apply a $\\tt{binary}$ quantizer with $\\tt{scale} = \\big\\lceil \\log_2(\\frac{\\sum |w|}{N}) \\big\\rceil$, where $w$ are the filter coefficients and $N$ is the number of coefficients, we will end up with the same filter $\\tt{binary}([-0.3, 0.2, 0.5, 0.15]) = \\tt{binary}([-0.5, 0.4, 0.1, 0.65]) = [-1,1,1,1] \\times 0.5$. In this case we are assuming the $\\tt{scale}$ is a power-of-2 number so that it can be efficiently implemented by a shift operation;\n",
    "\n",
    "- **more filters**: it is clear that quantization will drop information (just look at the example above) and deep quantization will drop more information, so to recover some of the boundary regions in layers that perform feature extraction, we may need to add more filters to the layer when we quantize it.\n",
    "\n",
    "We do not want to quantize the `softmax` layer, which is the last layer of the network. In `AutoQKeras`, you can specify the indexes that you want to perform quantization by specifying the corresponding index of the layer in `Keras`, i.e. if you can get the layer as `model.layers[i]` in `Keras`, `i` is the index of the layer.\n",
    "\n",
    "Finally, for data parallel distributed training, we should pass the strategy in `distribution_strategy` to `KerasTuner`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "height": 54
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 297,
     "status": "ok",
     "timestamp": 1591840405963,
     "user": {
      "displayName": "Claudionor Coelho",
      "photoUrl": "",
      "userId": "01084525977535968041"
     },
     "user_tz": 420
    },
    "id": "2-fyACb2dIAN",
    "outputId": "a180fa3f-8cc3-4f70-ce70-c05c28f88d1e"
   },
   "outputs": [],
   "source": [
    "run_config = {\n",
    "  \"output_dir\": tempfile.mkdtemp(),\n",
    "  \"goal\": goal,\n",
    "  \"quantization_config\": quantization_config,\n",
    "  \"learning_rate_optimizer\": False,\n",
    "  \"transfer_weights\": False,\n",
    "  \"mode\": \"random\",\n",
    "  \"seed\": 42,\n",
    "  \"limit\": limit,\n",
    "  \"tune_filters\": \"layer\",\n",
    "  \"tune_filters_exceptions\": \"^dense\",\n",
    "  \"distribution_strategy\": cur_strategy,\n",
    "  # first layer is input, layer two layers are softmax and flatten\n",
    "  \"layer_indexes\": range(1, len(model.layers) - 1),\n",
    "  \"max_trials\": 20\n",
    "}\n",
    "\n",
    "print(\"quantizing layers:\", [model.layers[i].name for i in run_config[\"layer_indexes\"]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "height": 1000
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 471192,
     "status": "ok",
     "timestamp": 1591840877167,
     "user": {
      "displayName": "Claudionor Coelho",
      "photoUrl": "",
      "userId": "01084525977535968041"
     },
     "user_tz": 420
    },
    "id": "MxlZFpa3fBv2",
    "outputId": "4d339846-1832-4a79-89b3-c9c4944dd47a"
   },
   "outputs": [],
   "source": [
    "autoqk = AutoQKeras(model, metrics=[\"acc\"], custom_objects=custom_objects, **run_config)\n",
    "autoqk.fit(x_train, y_train, validation_data=(x_test, y_test), batch_size=1024, epochs=20)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "LW_qN8-lOwL0"
   },
   "source": [
    "Now, let's see which model is the best model we got.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "height": 1000
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 3961,
     "status": "ok",
     "timestamp": 1591840881173,
     "user": {
      "displayName": "Claudionor Coelho",
      "photoUrl": "",
      "userId": "01084525977535968041"
     },
     "user_tz": 420
    },
    "id": "1L7KivAoffaL",
    "outputId": "f44b07a3-027d-4d69-9864-b3670815c407"
   },
   "outputs": [],
   "source": [
    "qmodel = autoqk.get_best_model()\n",
    "qmodel.save_weights(\"qmodel.h5\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "RB2xBRhJiwoh"
   },
   "source": [
    "We got here >90% reduction in energy when compared to 8-bit tensors and 32-bit accumulators. Remember that our original number was 3.3 uJ for fp32.  The end model has 11 nJ for the quantized model as opposed to 204 nJ for the 8-bit original quantized model. As these energy numbers are from high-level energy models, you should remember to consider the relations between them, and not the actual numbers."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Wy0zcqvQoBnb"
   },
   "source": [
    "Let's train this model to see how much accuracy we can get of it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "height": 1000
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 71353,
     "status": "ok",
     "timestamp": 1591840952535,
     "user": {
      "displayName": "Claudionor Coelho",
      "photoUrl": "",
      "userId": "01084525977535968041"
     },
     "user_tz": 420
    },
    "id": "_ipZSEfgoGdb",
    "outputId": "b184269d-1161-417a-e1ae-e852dc451561"
   },
   "outputs": [],
   "source": [
    "qmodel.load_weights(\"qmodel.h5\")\n",
    "with cur_strategy.scope():\n",
    "  optimizer = Adam(lr=0.02)\n",
    "  qmodel.compile(optimizer=optimizer, loss=\"categorical_crossentropy\", metrics=[\"acc\"])\n",
    "  qmodel.fit(x_train, y_train, epochs=200, batch_size=4096, validation_data=(x_test, y_test))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Fr95jcPROz7p"
   },
   "source": [
    "One of problems of trying to quantize the whole thing in one shot is that we may end up with too many choices to make, which will make the entire search space very high. In order to reduce the search space, `AutoQKeras` has two methods to enable users to cope with the explosion of choices."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "9zc7ZrnbPIJA"
   },
   "source": [
    "## Grouping Layers to Use the Same Choice\n",
    "\n",
    "In this case, we can provide regular expressions to `limit` to specify layer names that should be grouped together. In our example, suppose we want to group  convolution layers (except the first one) and all activations except the last one to use the same quantization.\n",
    "\n",
    "For the first convolution layer, we want to limit the quantization types to fewer choices as the input is already an 8-bit number.  The last activation will be fed to a feature classifier layer, so we may leave it with more bits. Because our `dense` is actually a `Conv2D` operation, we will enable 8-bits for the weights by layer name. \n",
    "\n",
    "We first need to look at the names of the layers for this. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "height": 428
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 301,
     "status": "ok",
     "timestamp": 1591840952867,
     "user": {
      "displayName": "Claudionor Coelho",
      "photoUrl": "",
      "userId": "01084525977535968041"
     },
     "user_tz": 420
    },
    "id": "w-d8nhG0pJF0",
    "outputId": "6529b630-f382-4e2a-94ef-ba3d9e3f875c"
   },
   "outputs": [],
   "source": [
    "pprint.pprint([layer.name for layer in model.layers])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "32Enp890pU_4"
   },
   "source": [
    "Convolution layers for `mnist` have names specified as `conv2d_[01234]`. Activation layers have names specified as `act_[01234]`. So, we can create the following regular expressions to reduce the search space in our model.\n",
    "\n",
    "Please note that layer class names always select different quantizers, so the user needs to specify a pattern for layer names if he/she wants to use the same quantization for the group of layers.\n",
    "\n",
    "You can see here another feature of the limit. You can specify the maximum number of bits, or cherry pick which quantizers you want to try for a specific layer if instead of the maximum number of bits you specify a list of quantizers fron `quantization_config`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Y5XItp95PHW6"
   },
   "outputs": [],
   "source": [
    "limit = {\n",
    "    \"Dense\": [8, 8, 4],\n",
    "    \"Conv2D\": [4, 8, 4],\n",
    "    \"DepthwiseConv2D\": [4, 8, 4],\n",
    "    \"Activation\": [4],\n",
    "    \"BatchNormalization\": [],\n",
    "\n",
    "    \"^conv2d_0$\": [\n",
    "                   [\"binary\", \"ternary\", \"quantized_bits(2,1,1,alpha=1.0)\"],\n",
    "                   8, 4\n",
    "    ],\n",
    "    \"^conv2d_[1234]$\": [4, 8, 4],\n",
    "    \"^act_[0123]$\": [4],\n",
    "    \"^act_4$\": [8],\n",
    "    \"^dense$\": [8, 8, 4]\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "EJs1L-jIie7w"
   },
   "outputs": [],
   "source": [
    "run_config = {\n",
    "  \"output_dir\": tempfile.mkdtemp(),\n",
    "  \"goal\": goal,\n",
    "  \"quantization_config\": quantization_config,\n",
    "  \"learning_rate_optimizer\": False,\n",
    "  \"transfer_weights\": False,\n",
    "  \"mode\": \"random\",\n",
    "  \"seed\": 42,\n",
    "  \"limit\": limit,\n",
    "  \"tune_filters\": \"layer\",\n",
    "  \"tune_filters_exceptions\": \"^dense\",\n",
    "  \"distribution_strategy\": cur_strategy,\n",
    "  \"layer_indexes\": range(1, len(model.layers) - 1),\n",
    "  \"max_trials\": 40\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "height": 1000
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 993665,
     "status": "ok",
     "timestamp": 1591841947161,
     "user": {
      "displayName": "Claudionor Coelho",
      "photoUrl": "",
      "userId": "01084525977535968041"
     },
     "user_tz": 420
    },
    "id": "c7eSwXyijhzc",
    "outputId": "6c76a21f-cbb3-4bc5-b899-b02c28821b78"
   },
   "outputs": [],
   "source": [
    "autoqk = AutoQKeras(model, metrics=[\"acc\"], custom_objects=custom_objects, **run_config)\n",
    "autoqk.fit(x_train, y_train, validation_data=(x_test, y_test), batch_size=1024, epochs=20)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "7sYp8Z2pnLi1"
   },
   "source": [
    "Let's see the reduction now."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "height": 1000
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 7109,
     "status": "ok",
     "timestamp": 1591841954308,
     "user": {
      "displayName": "Claudionor Coelho",
      "photoUrl": "",
      "userId": "01084525977535968041"
     },
     "user_tz": 420
    },
    "id": "yj826gNhjsfK",
    "outputId": "2e7f17d7-794e-44f6-d23a-452759727a53"
   },
   "outputs": [],
   "source": [
    "qmodel = autoqk.get_best_model()\n",
    "qmodel.save_weights(\"qmodel.h5\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "eXMcqxLAnY8t"
   },
   "source": [
    "Let's train this model for more time to see how much we can get in accuracy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "height": 1000
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 68145,
     "status": "ok",
     "timestamp": 1591842022471,
     "user": {
      "displayName": "Claudionor Coelho",
      "photoUrl": "",
      "userId": "01084525977535968041"
     },
     "user_tz": 420
    },
    "id": "qpT8QgkJnQPa",
    "outputId": "61e711db-6187-4047-dae8-9ce2d093f56c"
   },
   "outputs": [],
   "source": [
    "qmodel.load_weights(\"qmodel.h5\")\n",
    "with cur_strategy.scope():\n",
    "  optimizer = Adam(lr=0.02)\n",
    "  qmodel.compile(optimizer=optimizer, loss=\"categorical_crossentropy\", metrics=[\"acc\"])\n",
    "  qmodel.fit(x_train, y_train, epochs=200, batch_size=4096, validation_data=(x_test, y_test))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "gAV6Kw0QoODq"
   },
   "source": [
    "## Quantization by Blocks\n",
    "\n",
    "In the previous section, we enforced that all decisions were the same in order to reduce the number of options to quantize a model. \n",
    "\n",
    "Another approach is still to allow models to have each block of layers to makde their own choice, but quantizing the blocks sequentially, either from inputs to outputs, or by quantizing higher energy blocks first.\n",
    "\n",
    "The rationale for this method is that if we quantize the blocks one by one, and assuming that each block has $N$ choices, and $B$ blocks, we end up trying $N B$ options, instead of $N^B$ choices.  The reader should note that this is an approximation as there is no guarantee that we will obtain the best quantization possible.\n",
    "\n",
    "Should you do sequential from inputs to outputs or starting from the block that has the highest impact?\n",
    "\n",
    "If you have a network like ResNet, and if you want to do filter tuning, you need to block the layers by the resnet definition of a block, i.e. including full identity or convolutional blocks, and quantize the model from inputs to outputs, so that you can preserve at each stage the number of channels for the residual block. \n",
    "\n",
    "In order to perform quantization by blocks, you need to specify two other parameters in our `run_config`. `blocks` is a list of regular expressions of the groups you want to quantize. If a layer does not match the block pattern, it will not be quantized.  `schedule_block` specifies the mode for block quantization scheduling. It can be `sequential` or `cost` if you want to schedule first the blocks by decreasing cost size (energy or bits).\n",
    "\n",
    "In this model, there are a few optimizations that we perform automatically. First, we dynamically reduce the learning rate of the blocks that we have already quantized as setting them to not-trainable does not seem to work, so we still allow them to train, but at a slower pace. In addition, we try to dynamically adjust the learning rate for the layer we are trying to quantize as opposed to the learning rate of the unquantized layers. Finally, we transfer the weights of the models we have already quantized whenever we can do (if the shapes remain the same). \n",
    "\n",
    "Regardless on how we schedule the operations, we amortize the nubmer of trials for the cost of the block (energy or bits with respect to the total energy or number of bits of the network).\n",
    "\n",
    "Instead of invoking `AutoQKeras` now, we will invoke `AutoQKeras` scheduler."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "NUz4A6SKnhUf"
   },
   "outputs": [],
   "source": [
    "run_config = {\n",
    "  \"output_dir\": tempfile.mkdtemp(),\n",
    "  \"goal\": goal,\n",
    "  \"quantization_config\": quantization_config,\n",
    "  \"learning_rate_optimizer\": False,\n",
    "  \"transfer_weights\": False,\n",
    "  \"mode\": \"random\",\n",
    "  \"seed\": 42,\n",
    "  \"limit\": limit,\n",
    "  \"tune_filters\": \"layer\",\n",
    "  \"tune_filters_exceptions\": \"^dense\",\n",
    "  \"distribution_strategy\": cur_strategy,\n",
    "  \"layer_indexes\": range(1, len(model.layers) - 1),\n",
    "  \"max_trials\": 40,\n",
    "\n",
    "  \"blocks\": [\n",
    "    \"^.*_0$\",\n",
    "    \"^.*_1$\",\n",
    "    \"^.*_2$\",\n",
    "    \"^.*_3$\",\n",
    "    \"^.*_4$\",\n",
    "    \"^dense\"\n",
    "  ],\n",
    "  \"schedule_block\": \"cost\"\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "JWJiZZ9vsORJ"
   },
   "source": [
    "Because specifying regular expressions is error prone, we recommend that you first try to run `AutoQKerasScheduler` in debug mode to print the blocks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "height": 737
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 395,
     "status": "ok",
     "timestamp": 1591842023212,
     "user": {
      "displayName": "Claudionor Coelho",
      "photoUrl": "",
      "userId": "01084525977535968041"
     },
     "user_tz": 420
    },
    "id": "uSOxKQGwsqf2",
    "outputId": "18647e4f-ef7a-4c6a-aeb8-0c9c2039fdbb"
   },
   "outputs": [],
   "source": [
    "pprint.pprint([layer.name for layer in model.layers])\n",
    "autoqk = AutoQKerasScheduler(model, metrics=[\"acc\"], custom_objects=custom_objects, debug=True, **run_config)\n",
    "autoqk.fit(x_train, y_train, validation_data=(x_test, y_test), batch_size=1024, epochs=20)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "TQPUKPZhC_SI"
   },
   "source": [
    "All blocks seem to be fine. Let's find the best quantization now."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "height": 1000
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 1938883,
     "status": "ok",
     "timestamp": 1591843962106,
     "user": {
      "displayName": "Claudionor Coelho",
      "photoUrl": "",
      "userId": "01084525977535968041"
     },
     "user_tz": 420
    },
    "id": "sXt-cRKvDEaL",
    "outputId": "36db3217-86ff-4425-ee12-f637a4fc1841"
   },
   "outputs": [],
   "source": [
    "autoqk = AutoQKerasScheduler(model, metrics=[\"acc\"], custom_objects=custom_objects, **run_config)\n",
    "autoqk.fit(x_train, y_train, validation_data=(x_test, y_test), batch_size=1024, epochs=20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "height": 291
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 396,
     "status": "ok",
     "timestamp": 1591843962540,
     "user": {
      "displayName": "Claudionor Coelho",
      "photoUrl": "",
      "userId": "01084525977535968041"
     },
     "user_tz": 420
    },
    "id": "ArdGbsXFDK-I",
    "outputId": "43730cd5-93fc-4838-c49a-1f3f4151fa54"
   },
   "outputs": [],
   "source": [
    "qmodel = autoqk.get_best_model()\n",
    "qmodel.save_weights(\"qmodel.h5\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "height": 1000
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 69779,
     "status": "ok",
     "timestamp": 1591844032332,
     "user": {
      "displayName": "Claudionor Coelho",
      "photoUrl": "",
      "userId": "01084525977535968041"
     },
     "user_tz": 420
    },
    "id": "RHGb6YHFEgtV",
    "outputId": "5578ce49-1ee9-4063-deab-1b3db9f4b66b"
   },
   "outputs": [],
   "source": [
    "qmodel.load_weights(\"qmodel.h5\")\n",
    "with cur_strategy.scope():\n",
    "  optimizer = Adam(lr=0.02)\n",
    "  qmodel.compile(optimizer=optimizer, loss=\"categorical_crossentropy\", metrics=[\"acc\"])\n",
    "  qmodel.fit(x_train, y_train, epochs=200, batch_size=4096, validation_data=(x_test, y_test))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "fJCkMdAcjnoh"
   },
   "source": [
    "Perfect! You have learned how to perform automatic quantization using AutoQKeras with QKeras."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "last_runtime": {
    "build_target": "//learning/deepmind/dm_python:dm_notebook3_tpu",
    "kind": "private"
   },
   "name": "AutoQKeras.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}