test_edgetpu_pixellink.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Demonstrate construction and quantization of a model for execution on the EdgeTPU\n",
    "\n",
    "This MNIST example should be helpful:  \n",
    "https://www.tensorflow.org/lite/performance/post_training_integer_quant  \n",
    "\n",
    "This webpage, and links at the bottom of the page, may also be helpful:  \n",
    "https://www.tensorflow.org/lite/performance/post_training_quantization\n",
    "\n",
    "## See of some of these comments help too:  \n",
    "https://github.com/google-coral/edgetpu/issues/13"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "from tensorflow.keras.models import Model\n",
    "from tensorflow.keras.layers import Input, Conv2D, MaxPooling2D, SeparableConv2D, ReLU\n",
    "import time\n",
    "\n",
    "# %matplotlib notebook\n",
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Print out some system information\n",
    "import subprocess\n",
    "\n",
    "print('Linux kernel version:')\n",
    "print('$ uname -r')\n",
    "result = subprocess.run(['uname', '-r'], stdout=subprocess.PIPE)\n",
    "print(result.stdout.decode('utf-8'))\n",
    "\n",
    "print('Linux release:')\n",
    "print('$ lsb_release -a')\n",
    "result = subprocess.run(['lsb_release', '-a'], stdout=subprocess.PIPE)\n",
    "print(result.stdout.decode('utf-8'))\n",
    "\n",
    "print('Tensorflow python module version')\n",
    "print(tf.__version__)\n",
    "print('')\n",
    "\n",
    "print('Edge TPU python module version:')\n",
    "import edgetpu\n",
    "print(edgetpu.__version__)\n",
    "print('')\n",
    "\n",
    "print('Edge TPU compiler version:')\n",
    "result = subprocess.run(['edgetpu_compiler', '--version'], stdout=subprocess.PIPE)\n",
    "print(result.stdout.decode('utf-8'))\n",
    "print('')\n",
    "\n",
    "print('Edge TPU runtime file:')\n",
    "result = subprocess.check_output(\"dpkg -l | grep libedgetpu\", shell=True)\n",
    "print(result.decode(\"utf-8\"))\n",
    "\n",
    "print('Edge TPU runtime version:')\n",
    "import edgetpu.basic.edgetpu_utils\n",
    "print(edgetpu.basic.edgetpu_utils.GetRuntimeVersion())\n",
    "\n",
    "# import tflite_runtime.interpreter as tflite\n",
    "import tflite_runtime\n",
    "print('tflite_runtime version:')\n",
    "print(tflite_runtime.__version__)\n",
    "print('')\n",
    "\n",
    "print('Paths of available Edge TPU devices, if any:')\n",
    "devices = edgetpu.basic.edgetpu_utils.ListEdgeTpuPaths(edgetpu.basic.edgetpu_utils.EDGE_TPU_STATE_NONE)\n",
    "print(devices)\n",
    "# Set identity of edge device to use, if any\n",
    "if len(devices) > 0:\n",
    "    # Use the first device in the list\n",
    "    if devices[0].startswith('/dev/apex'):\n",
    "        target_device = 'pci'\n",
    "    else:\n",
    "        # Assuming device is on USB bus\n",
    "        target_device = 'usb'\n",
    "    print('Using delegate device: \"%s\"' % (target_device))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if tf.__version__.startswith('1.15'):\n",
    "    # This prevents some errors that otherwise occur when converting the model with TF 1.15...\n",
    "    tf.enable_eager_execution() # Only if TF is version 1.15    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "## Build a PixelLink model\n",
    "\n",
    "import model_architectures as model_archs\n",
    "\n",
    "gt_res='2s'  # resolution of ground truth bitmaps. '2s' or '4s'\n",
    "# image_shape = (416, 416, 3) # works for: n_channels = [32, 64, 128, 256, 256, 512, 512]\n",
    "# image_shape = (448, 448, 3) # output has artifact for: n_channels = [32, 64, 128, 256, 256, 512, 512]\n",
    "# image_shape = (512, 512, 3) # freezes for: n_channels = [32, 64, 128, 256, 256, 512, 512]\n",
    "image_shape = (512, 512, 3) # works for: n_channels = [16, 32, 64, 128, 128, 256, 256]\n",
    "n_bitmaps = 5\n",
    "\n",
    "def representative_dataset_gen():\n",
    "    num_calibration_images = 100\n",
    "    for i in range(num_calibration_images):\n",
    "        image = tf.random.uniform([1] + list(image_shape))\n",
    "        yield [image]\n",
    "\n",
    "# n_channels = [64, 128, 256, 512, 512, 1024, 1024], # PixelLink\n",
    "# n_channels = [32, 64, 128, 256, 256, 512, 512] # Primary model, works for up to 416x416 input with sep convs\n",
    "n_channels = [16, 32, 64, 128, 128, 256, 256]\n",
    "# n_channels = [8, 16, 32, 64, 64, 128, 128]\n",
    "kernel_sizes = [3, 3, 3, 3, 3, 1, 1]\n",
    "n_convs = [2, 2, 3, 3, 3, 1, 1]\n",
    "pool = [True, True, True, True, False, False, False]\n",
    "build_model = lambda : model_archs.build_modular(image_shape, n_gt_chans=n_bitmaps,\n",
    "                                                 resolution=gt_res, separable_conv=True,\n",
    "                                                 batchnorm='batchnorm', n_channels=n_channels,\n",
    "                                                 kernel_sizes=kernel_sizes, n_convs=n_convs, pool=pool)\n",
    "model = build_model()\n",
    "\n",
    "# model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Convert to tensorflow lite model and save...\n",
    "if tf.__version__.startswith('2.'):\n",
    "    converter = tf.lite.TFLiteConverter.from_keras_model(model) # TF2.0\n",
    "elif tf.__version__.startswith('1.15'):\n",
    "    model.save('model_keras', include_optimizer=False) # TF1.15\n",
    "    converter = tf.lite.TFLiteConverter.from_keras_model_file('model_keras') # TF1.15\n",
    "else:\n",
    "    raise ValueError('Unhandled TensorFlow version.')\n",
    "    \n",
    "converter.optimizations = [tf.lite.Optimize.DEFAULT]\n",
    "converter.representative_dataset = tf.lite.RepresentativeDataset(representative_dataset_gen)\n",
    "converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8] # For EdgeTPU, no float ops allowed\n",
    "\n",
    "tflite_model = converter.convert()\n",
    "open('model.tflite', 'wb').write(tflite_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "##================================================================================================\n",
    "## Set variable below and comment out lines in cell below if not using Edge TPU model and hardware\n",
    "##================================================================================================\n",
    "# True:  Use EdgeTPU model and process on the Edge TPU (assumes one is available)\n",
    "# False: Use TFLite model and process on CPU\n",
    "use_edgetpu = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "%%bash\n",
    "## Compile model for edge TPU\n",
    "# Note that the output file name has '_edgetpu' appended to the root filename of the input TFLite model.\n",
    "edgetpu_compiler --min_runtime_version 13 --show_operations 'model.tflite'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Load TFLite model and allocate tensors.\n",
    "if use_edgetpu:\n",
    "    # If using interpreter from tflite_runtime package\n",
    "    from tflite_runtime.interpreter import load_delegate\n",
    "    from tflite_runtime.interpreter import Interpreter\n",
    "    interpreter = Interpreter(model_path='model_edgetpu.tflite',\n",
    "                              model_content=None,\n",
    "                              experimental_delegates=[load_delegate('libedgetpu.so.1.0',\n",
    "                                                                    {'device': target_device})])\n",
    "    \n",
    "#     # If using interpreter from full TensorFlow package...\n",
    "#     from tensorflow.lite.python.interpreter import load_delegate\n",
    "#     interpreter = tf.lite.Interpreter(model_path='model_edgetpu.tflite',\n",
    "#                                       experimental_delegates=[load_delegate('libedgetpu.so.1.0',\n",
    "#                                                                             {'device': target_device})])\n",
    "else:\n",
    "    interpreter = tf.lite.Interpreter(model_path='model.tflite')\n",
    "\n",
    "interpreter.allocate_tensors()\n",
    "\n",
    "# Get input and output tensors details.\n",
    "input_details = interpreter.get_input_details()\n",
    "output_details = interpreter.get_output_details()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Put some random data through the model and show results\n",
    "# Create a batch of images\n",
    "batch_size = 10\n",
    "# image = tf.random.normal([batch_size] + list(image_shape))\n",
    "image = tf.random.uniform([batch_size] + list(image_shape))\n",
    "# image = tf.random.uniform([batch_size] + list(image_shape))*256\n",
    "# image = tf.cast(image, tf.dtypes.uint8)\n",
    "\n",
    "# Process the image with the network model\n",
    "t_all = time.time()\n",
    "t_individual = np.zeros(batch_size)\n",
    "for i_im in range(batch_size):\n",
    "    t_one = time.time()\n",
    "    # Set input tensor and invoke model\n",
    "    interpreter.set_tensor(input_details[0]['index'], image[i_im:i_im+1])\n",
    "    interpreter.invoke()   # Can be slow if running on CPU\n",
    "\n",
    "    # The function `get_tensor()` returns a copy of the tensor data.\n",
    "    # Use `tensor()` in order to get a pointer to the tensor.\n",
    "    model_output = interpreter.get_tensor(output_details[0]['index'])\n",
    "    t_individual[i_im] = time.time() - t_one\n",
    "print('Model processing took %f seconds.' % (time.time() - t_all))\n",
    "print('Individual image processing times:')\n",
    "print(t_individual)\n",
    "\n",
    "# Plot results for first channel of input and output, of the first\n",
    "# image in the batch.\n",
    "in_chan0 = image[0, :, :, 0]\n",
    "out_chan0 = model_output[0, :, :, 0]\n",
    "plt.figure(figsize=(16, 8))\n",
    "# clim = [-4, 4]\n",
    "\n",
    "plt.subplot(1,2,1)\n",
    "plt.imshow(in_chan0, aspect='equal')\n",
    "plt.title('Input min chan0 value: %f' % (tf.reduce_min(in_chan0)))\n",
    "# plt.clim(clim)\n",
    "plt.colorbar()\n",
    "\n",
    "plt.subplot(1,2,2)\n",
    "plt.imshow(out_chan0, aspect='equal')\n",
    "plt.title('Output min chan0 value: %f' % (tf.reduce_min(out_chan0)))\n",
    "# plt.clim(clim)\n",
    "plt.colorbar()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "env_tensorflow2",
   "language": "python",
   "name": "env_tensorflow2"
  },
  "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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}