quantization_edgetpu.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Quantization and Edge TPU",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r4HQauY6RivH",
        "colab_type": "text"
      },
      "source": [
        "**Post-training quantization and quantized inference**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "USkKyQVHzvbr",
        "colab_type": "text"
      },
      "source": [
        "Post-training quantization is a conversion technique that can **reduce model size** while also **improving CPU and hardware accelerator latency**, with little **degradation** in model **accuracy.** You can perform these techniques using an already-trained **float** TensorFlow model when you convert it to TensorFlow Lite format. Once a model is **fully quantized**, you can deploy it to platforms like **coral tpu** or run it using **nnapi** **delegates**."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wK8kRTkTSIt7",
        "colab_type": "text"
      },
      "source": [
        " Install latest **tf-nightly** for quantization and tflite inference."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gWPepcH3AmRp",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!pip install tf-nightly"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "aPcnVG0MjBux",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!unzip /content/drive/My\\ Drive/portrait256/portrait256.zip"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IoSc7RJKA0_Q",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BS4afMkHWXkT",
        "colab_type": "text"
      },
      "source": [
        "The following diagram shows the various methods for **post-training quantization** available in tensorflow-lite.\n",
        "\n",
        "![TfliteQuantizatiom](https://www.tensorflow.org/lite/performance/images/optimization.jpg)\n",
        "\n",
        "Here we will use **full integer quantization**(INT8) with a representative dataset. Use the **trainig data-set** as the representative dataset for this purpose."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uHKPHxqmAxmZ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!unzip /content/drive/My\\ Drive/portrait_mix/EG1800.zip"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "51CWZpJ7G8yT",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Import libraries\n",
        "import tensorflow as tf\n",
        "import cv2, sys, time\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "from PIL import Image\n",
        "import timeit\n",
        "from matplotlib import pyplot as plt"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "T5lGb-v5Ltfx",
        "colab_type": "text"
      },
      "source": [
        "To convert a float32 model to quantized int8 format, there are two approaches viz. **quantizaton aware training** and **post training quantization**. Here we use the latter approach of quantization, to convert our model to **INT8** format. We need to quantize both the weights and the corresponding layers to make use of accelerators like **Coral TPU or NNAPI**. We will use the training dataset as a **representative dataset** to measure the dynamic range of activations and inputs.\n",
        "\n",
        "**Note:** TF 2.0/Keras currently does not support quantization aware training."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LoaAcChbAWE0",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Configure dataset directories\n",
        "IMGDS='/content/EG1800/imgs'\n",
        "MSKDS='/content/EG1800/msks'\n",
        "\n",
        "# Use image data generator to load the dataset\n",
        "image_generator = tf.keras.preprocessing.image.ImageDataGenerator()\n",
        "seed=7\n",
        "\n",
        "train_image_generator = image_generator.flow_from_directory(\n",
        "    IMGDS,\n",
        "    batch_size=8,\n",
        "    shuffle=True,\n",
        "    target_size=(224, 224),\n",
        "    color_mode=\"rgb\",\n",
        "    class_mode=None,\n",
        "    seed=seed)\n",
        "\n",
        "train_mask_generator = image_generator.flow_from_directory(\n",
        "    MSKDS,\n",
        "    batch_size=8,\n",
        "    shuffle=True,\n",
        "    target_size=(224, 224),\n",
        "    color_mode=\"grayscale\",\n",
        "    class_mode=None,\n",
        "    seed=seed)\n",
        "\n",
        "# Normalize the input image\n",
        "def normalize(imgOri, scale=1, mean=[103.94, 116.78, 123.68], val=[0.017, 0.017, 0.017]):\n",
        "    img = np.array(imgOri.copy(), np.float32)/scale\n",
        "    return (img - mean) * val"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QtpsLHTGOw_O",
        "colab_type": "text"
      },
      "source": [
        "Use the keras **data generator** to load the dataset and perform the data augmentaion(like training) to generate the inputs for the quantization process. Configure the tflite converter for **full integer quantization** and convert the model to **INT8 tflite** format. Now the saved model will have **INT8 Range: -128 to 127**."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4bsHaxQcAdPm",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Load the dataset with data augmentations\n",
        "def representative_dataset_gen():\n",
        "  for _ in range(100):\n",
        "    imgs=train_image_generator.next()\n",
        "    msks=train_mask_generator.next()\n",
        "    input=np.zeros(shape=(8,224,224,4), dtype=np.float32)\n",
        "    for i in range(len(imgs)):\n",
        "      img=normalize(cv2.cvtColor(imgs[i],cv2.COLOR_RGB2BGR))\n",
        "      msk=cv2.normalize(cv2.blur(msks[i], (5,5)), 0, 1, cv2.NORM_MINMAX) # partial augmentation\n",
        "      input=np.float32(np.dstack([img,msk]).reshape(1,224,224,4))\n",
        "\n",
        "      yield [input]\n",
        "\n",
        "# Convert the model to INT8 format\n",
        "converter = tf.compat.v1.lite.TFLiteConverter.from_keras_model_file('portrait_video_flattened.h5')\n",
        "converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]\n",
        "converter.inference_input_type = tf.int8\n",
        "converter.inference_output_type = tf.int8\n",
        "converter.optimizations = [tf.lite.Optimize.DEFAULT]\n",
        "converter.representative_dataset = representative_dataset_gen\n",
        "\n",
        "# Save the quantized model\n",
        "tflite_quant_model = converter.convert()\n",
        "open(\"/content/portrait_video_quant.tflite\", \"wb\").write(tflite_quant_model)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZtLiLK5gP3fO",
        "colab_type": "text"
      },
      "source": [
        "Test the new quantized model on a new **video file**, with the help of opencv. Make sure to convert the float input to **INT8 range** before inference and INT8 output to **float32** after quantized inference. Save the outputs using opencv video writer and **compare** the output with the original results."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wxzpZ_a7AhtM",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Normalize the input image\n",
        "def normalize(imgOri, scale=1, mean=[103.94, 116.78, 123.68], val=[0.017, 0.017, 0.017]):\n",
        "    img = np.array(imgOri.copy(), np.float32)/scale\n",
        "    norm=(img - mean) * val\n",
        "    out=cv2.normalize(src=norm, dst=None, alpha=-128, beta=127, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_8S)\n",
        "    return out\n",
        " \n",
        "# Alpha blend frame with background\n",
        "def blend(frame, alpha):\n",
        "        background = np.zeros(frame.shape) + [255, 255, 255]\n",
        "        alphargb = cv2.cvtColor(alpha, cv2.COLOR_GRAY2BGR)\n",
        "        result = np.uint8(frame * alphargb + background * (1-alphargb))\n",
        "        return frame, alphargb*255, result\n",
        "\n",
        "# Initialize tflite-interpreter\n",
        "interpreter = tf.lite.Interpreter(model_path=\"portrait_video_fquant.tflite\") # Use 'tf.lite' on recent tf versions\n",
        "interpreter.allocate_tensors()\n",
        "input_details = interpreter.get_input_details()\n",
        "output_details = interpreter.get_output_details()\n",
        "input_shape = input_details[0]['shape'][1:3]\n",
        "\n",
        "\n",
        "# Initialize video capturer\n",
        "videofile = 'portrait_lady.mp4'\n",
        "cap = cv2.VideoCapture(videofile)  \n",
        "size = (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)))  \n",
        "videoWriter = cv2.VideoWriter('result.mp4', cv2.VideoWriter_fourcc(*'MJPG'), 20, size)  \n",
        " \n",
        "# Initialize frame counter \n",
        "cnt = 1\n",
        "\n",
        "while True:\n",
        "       \n",
        "    # Read the BGR frames \n",
        "    success, frame = cap.read()\n",
        "    if not(success):\n",
        "      break\n",
        "    image=Image.fromarray(frame)\n",
        "    \n",
        "    \n",
        "    # Resize the image\n",
        "    image= image.resize(input_shape, Image.ANTIALIAS)\n",
        "    image=np.asarray(image)\n",
        "\n",
        "    \n",
        "    # Normalize the input\n",
        "    image = normalize(image)\n",
        "  \n",
        "    # Choose prior mask\n",
        "    if cnt == 1:\n",
        "        prior = np.full((224, 224, 1),-128,dtype=np.int8) # first frame\n",
        "    else:\n",
        "        prior = pred_video\n",
        "    \n",
        "    # Add prior as fourth channel\n",
        "    image=np.dstack([image,prior])\n",
        "    prepimg = image[np.newaxis, :, :, :]\n",
        "    \n",
        "\n",
        "    # Invoke interpreter for inference\n",
        "    interpreter.set_tensor(input_details[0]['index'], np.array(prepimg, dtype=np.int8))\n",
        "    interpreter.invoke()\n",
        "    outputs = interpreter.get_tensor(output_details[0]['index'])\n",
        "    outputs = outputs.reshape(224,224,1)\n",
        "  \n",
        "    # Save output to feed subsequent inputs\n",
        "    pred_video = outputs\n",
        "\n",
        "    # Process the output and perform alpha blending\n",
        "    outputs = outputs.astype('float32')  \n",
        "    outputs = cv2.resize(outputs, size)\n",
        "    outputs= cv2.normalize(src=outputs, dst=None, alpha=0.0, beta=1.0, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_32F)\n",
        "    _,_,outputs=blend(frame, outputs)\n",
        "        \n",
        "    # Write the output frame\n",
        "    videoWriter.write(outputs)\n",
        "    \n",
        "    # Print the frame count\n",
        "    cnt += 1\n",
        "    if cnt % 100 == 0:\n",
        "        print (\"cnt: \", cnt)\n",
        "\n",
        "# When everything done, release the capturer\n",
        "print(\"Conversion successfull !!!\")\n",
        "videoWriter.release()\n",
        "cap.release()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d9RZhlE-VKoV",
        "colab_type": "text"
      },
      "source": [
        "Older versions of **tensorflow**(1.x) does not seem to support conversions for quantized versions of **transpose convolutions** using the current approach. Compared to the original float model, the output with quantized model looks **less accurate**. Hopefully, the resut may be slighlty improved with **full data augmentation**  and/or **quantization aware training**. Unfortunately the fully quantized model could not be converted to **coral tpu** format due to unsupported layers and compatability issues; but it may work with the latest **NNAPI** delegate in android."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OUGXxqhTuaJI",
        "colab_type": "text"
      },
      "source": [
        "**Prisma-net: Quantization & TPU Conversion**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LyMioPdGwudO",
        "colab_type": "text"
      },
      "source": [
        "Let's quantize our old **prisma-net** model and see if it can be successfully converted to **TPU format**. Firtstly, intitialize the **data generator** for loading the dataset used for **training** prisma-net."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "R32wkf7ewa9S",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# prisma data generator\n",
        "IMGDS=\"/content/portrait256/images\";\n",
        "\n",
        "data_gen_args = dict(rescale=1./255,\n",
        "                     width_shift_range=0.2,\n",
        "                     height_shift_range=0.2,\n",
        "                     zoom_range=0.2,\n",
        "                     horizontal_flip=True,\n",
        "                     validation_split=0.2\n",
        "                    )\n",
        "\n",
        "image_datagen = tf.keras.preprocessing.image.ImageDataGenerator(**data_gen_args)\n",
        "\n",
        "# Provide the same seed and keyword arguments to the fit and flow methods\n",
        "seed = 1\n",
        "\n",
        "train_image_generator = image_datagen.flow_from_directory(\n",
        "    IMGDS,\n",
        "    batch_size=8,\n",
        "    shuffle=True,\n",
        "    subset='training',\n",
        "    color_mode=\"rgb\",\n",
        "    class_mode=None,\n",
        "    seed=seed)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y5vjSn2pxKTL",
        "colab_type": "text"
      },
      "source": [
        "Load the **keras model**, configure the **tflite converter** and start the quantization process using the **representative dataset**."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RF_LGDdnwgU5",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Load the dataset with data augmentations\n",
        "def representative_dataset_gen():\n",
        "  for _ in range(200):\n",
        "    imgs=train_image_generator.next()\n",
        "    for i in range(len(imgs)):\n",
        "      img=imgs[i]\n",
        "      input=np.float32(img).reshape(1,256,256,3)\n",
        "      yield [input]\n",
        "\n",
        "# Convert the model to INT8 format\n",
        "converter = tf.compat.v1.lite.TFLiteConverter.from_keras_model_file('prisma-net-15-0.08.hdf5')\n",
        "converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]\n",
        "converter.inference_input_type = tf.int8\n",
        "converter.inference_output_type = tf.int8\n",
        "converter.optimizations = [tf.lite.Optimize.DEFAULT]\n",
        "converter.representative_dataset = representative_dataset_gen\n",
        "\n",
        "# Save the quantized model\n",
        "tflite_quant_modell = converter.convert()\n",
        "open(\"/content/prisma_quant.tflite\", \"wb\").write(tflite_quant_modell)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SRY4A1Xs1dMR",
        "colab_type": "text"
      },
      "source": [
        "**Test** the quantized model on a portrait image"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ocyVtWCv1cRJ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Read the input image and normalize it\n",
        "img=Image.open('/content/kid8.jpg').resize((256,256),Image.BICUBIC)\n",
        "img=np.array(img)\n",
        "norm=np.int8(img-128) # Normalize to INT8\n",
        "prepimg=norm.reshape(1,256,256,3)\n",
        "\n",
        "# Initialize tflite-interpreter\n",
        "interpreter = tf.lite.Interpreter(model_path=\"prisma_quant.tflite\")\n",
        "interpreter.allocate_tensors()\n",
        "input_details = interpreter.get_input_details()\n",
        "output_details = interpreter.get_output_details()\n",
        "input_shape = input_details[0]['shape'][1:3]\n",
        "\n",
        "# Invoke interpreter for inference\n",
        "interpreter.set_tensor(input_details[0]['index'], np.array(prepimg, dtype=np.int8))\n",
        "interpreter.invoke()\n",
        "outputs = interpreter.get_tensor(output_details[0]['index'])\n",
        "outputs = outputs.reshape(256,256,1)\n",
        "\n",
        "# Plot the output results\n",
        "result=np.uint8(outputs+128) # Denormalize to UINT8\n",
        "result=np.uint8((result/255.0)*img) # Crop the image with mask\n",
        "plt.imshow(result)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UrVgdhvBx4Ww",
        "colab_type": "text"
      },
      "source": [
        "Install **coral edge-tpu compiler** on the host system."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fmkcR7AyrSqw",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key add -\n",
        "!echo \"deb https://packages.cloud.google.com/apt coral-edgetpu-stable main\" | sudo tee /etc/apt/sources.list.d/coral-edgetpu.list\n",
        "!sudo apt-get update\n",
        "!sudo apt-get install edgetpu"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "u6a9Aj5ax_2D",
        "colab_type": "text"
      },
      "source": [
        "Convert the quantized tflite to **TPU format**."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rzhWrkpmvqUx",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!edgetpu_compiler prisma_quant.tflite"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4mQvehbcyKzv",
        "colab_type": "text"
      },
      "source": [
        "\n",
        "Here is the **output** of the  conversion process\n",
        "\n",
        "```\n",
        "Edge TPU Compiler version 2.0.291256449\n",
        "\n",
        "Model compiled successfully in 757 ms.\n",
        "\n",
        "Input model: prisma_quant.tflite\n",
        "Input size: 1.15MiB\n",
        "Output model: prisma_quant_edgetpu.tflite\n",
        "Output size: 1.31MiB\n",
        "On-chip memory available for caching model parameters: 7.43MiB\n",
        "On-chip memory used for caching model parameters: 1.13MiB\n",
        "Off-chip memory used for streaming uncached model parameters: 64.00B\n",
        "Number of Edge TPU subgraphs: 1\n",
        "Total number of operations: 158\n",
        "Operation log: prisma_quant_edgetpu.log\n",
        "See the operation log file for individual operation details.\n",
        "```\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VrLxi6Ef0hSD",
        "colab_type": "text"
      },
      "source": [
        "Now to run your model on a coral TPU, we need to install the **Edge TPU Runtime and Tensorflow Lite Libraries**. Install the appropriate versions by following the **coral** documentation:-\n",
        "\n",
        "1. **Edge TPU Runtime**: https://coral.ai/docs/accelerator/get-started/#1-install-the-edge-tpu-runtime\n",
        "2. **Tensorflow Lite Library**: https://coral.ai/docs/accelerator/get-started/#2-install-the-tensorflow-lite-library\n",
        "\n",
        "Finally connect the **Edge TPU** to host PC(USB 3.0) and **verify** the installation by running the demo **classification** example provided in the documentaion. To run the **inference** on the model, we can use the same Tflite API (or Edge TPU Python API ). There is only one major **change** in the code compared to the normal tflite inference ie.**experimental_delegates** argument."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Juzl70ZY2siJ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import numpy as np\n",
        "from PIL import Image\n",
        "import tflite_runtime.interpreter as tflite\n",
        "from imageio import imsave\n",
        "import time\n",
        "\n",
        "# Read the input image and normalize it\n",
        "img=Image.open('kid8.jpg').resize((256,256),Image.BICUBIC)\n",
        "img=np.array(img)\n",
        "norm=np.int8(img-128) # Normalize to INT8\n",
        "prepimg=norm.reshape(1,256,256,3)\n",
        "\n",
        "# Initialize tflite-interpreter\n",
        "interpreter = tflite.Interpreter(model_path=\"prisma_quant_edgetpu.tflite\", experimental_delegates=[tflite.load_delegate('libedgetpu.so.1')]) # Use 'tf.lite' on recent tf versions\n",
        "interpreter.allocate_tensors()\n",
        "input_details = interpreter.get_input_details()\n",
        "output_details = interpreter.get_output_details()\n",
        "input_shape = input_details[0]['shape'][1:3]\n",
        "\n",
        "# Invoke interpreter for inference\n",
        "interpreter.set_tensor(input_details[0]['index'], np.array(prepimg, dtype=np.int8))\n",
        "\n",
        "# Mesure execution time\n",
        "for _ in range(5):\n",
        "    start = time.perf_counter()\n",
        "    interpreter.invoke()\n",
        "    inference_time = time.perf_counter() - start\n",
        "    outputs = interpreter.get_tensor(output_details[0]['index'])\n",
        "    print('%.1fms' % (inference_time * 1000))\n",
        "\n",
        "\n",
        "outputs = outputs.reshape(256,256,1)\n",
        "\n",
        "# Plot the output results\n",
        "result=np.uint8(outputs+128) # Denormalize to UINT8\n",
        "result=np.uint8((result/255.0)*img) # Crop the image with mask\n",
        "imsave(\"segout.png\",result)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EjvtMgmN6fqk",
        "colab_type": "text"
      },
      "source": [
        "**Results:-**\n",
        "```\n",
        "Tflite Quantized CPU: 4357.0 ms\n",
        "Tflite Quantized TPU:   12.2 ms\n",
        "Tflite Float32 CPU:    500.0 ms\n",
        "```\n",
        "The **CPU**(i7-3632QM CPU @ 2.20GHz) might be using a **single core** for inference. But even if we include other possible overheads, this **40x speed-up** seems to be worth the effort. Besides, it consumes **20 times less power** than CPU."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kLgKzTzK3mhd",
        "colab_type": "text"
      },
      "source": [
        "**Note:** The official documenation only recommends **TF 1.x** with Edge TPU. There is still a confusion regarding input types i.e **UINT8 vs INT8**. However the model was **successfully converted** and it seems to give **correct ouput** with the device during inference."
      ]
    }
  ]
}