{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wJcYs_ERTnnI"
      },
      "source": [
        "##### Copyright 2021 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "HMUDt0CiUJk9"
      },
      "outputs": [],
      "source": [
        "#@title 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": {
        "id": "77z2OchJTk0l"
      },
      "source": [
        "# Migrating your TFLite code to TF2\n",
        "\n",
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://www.tensorflow.org/guide/migrate/tflite\">\n",
        "    <img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />\n",
        "    View on TensorFlow.org</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs/blob/master/site/en/guide/migrate/tflite.ipynb\">\n",
        "    <img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />\n",
        "    Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/tensorflow/docs/blob/master/site/en/guide/migrate/tflite.ipynb\">\n",
        "    <img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />\n",
        "    View source on GitHub</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://storage.googleapis.com/tensorflow_docs/docs/site/en/guide/migrate/tflite.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\" />Download notebook</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "meUTrR4I6m1C"
      },
      "source": [
        "[TensorFlow Lite](https://www.tensorflow.org/lite/guide) (TFLite) is a set of tools that helps developers run ML inference on-device (mobile, embedded, and IoT devices). The [TFLite converter](https://www.tensorflow.org/lite/convert)  is one such tool that converts existing TF models into an optimized TFLite model format that can be efficiently run on-device.\n",
        "\n",
        "In this doc, you'll learn what changes you need to make to your TF to TFLite  conversion code, followed by a few examples that do the same.\n",
        "\n",
        "\n",
        "## Changes to your TF to TFLite  conversion code\n",
        "\n",
        "* If you're using a legacy TF1 model format (such as Keras file, frozen GraphDef, checkpoints, tf.Session), update it to TF1/TF2 SavedModel and use the TF2 converter API `tf.lite.TFLiteConverter.from_saved_model(...)` to convert it to a TFLite model (refer to Table 1).\n",
        "\n",
        "* Update the converter API flags (refer to Table 2).\n",
        "* Remove legacy APIs such as `tf.lite.constants`. (eg: Replace `tf.lite.constants.INT8` with `tf.int8`)\n",
        "\n",
        "// Table 1 // TFLite Python Converter API Update\n",
        "\n",
        "TF1 API | TF2 API |\n",
        "--- | --- |\n",
        "`tf.lite.TFLiteConverter.from_saved_model('saved_model/',..)` | *supported* |\n",
        "`tf.lite.TFLiteConverter.from_keras_model_file('model.h5',..)` | *removed (update to SavedModel format)* |\n",
        "`tf.lite.TFLiteConverter.from_frozen_graph('model.pb',..)` | *removed (update to SavedModel format)* |\n",
        "`tf.lite.TFLiteConverter.from_session(sess,...)` | *removed (update to SavedModel format)* |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Rf75rjeedigq"
      },
      "source": [
        "<style>\n",
        "  .table {margin-left: 0 !important;}\n",
        "</style>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XbVlZNizW1-Y"
      },
      "source": [
        "// Table 2 // TFLite Python Converter API Flags Update\n",
        "\n",
        "TF1 API | TF2 API |\n",
        "--- | --- |\n",
        "`allow_custom_ops`<br>`optimizations`<br> `representative_dataset`<br>`target_spec` <br>`inference_input_type`<br>`inference_output_type`<br>`experimental_new_converter`<br> `experimental_new_quantizer` | *supported* <br><br><br><br><br><br><br><br> |\n",
        "`input_tensors`<br>`output_tensors`<br>`input_arrays_with_shape`<br>`output_arrays`<br>`experimental_debug_info_func`| *removed (unsupported converter API arguments)*<br><br><br><br><br> |\n",
        "`change_concat_input_ranges`<br>`default_ranges_stats`<br>`get_input_arrays()`<br>`inference_type`<br>`quantized_input_stats`<br> `reorder_across_fake_quant` | *removed (unsupported quantization workflows)*<br><br><br><br><br><br> |\n",
        "`conversion_summary_dir`<br>`dump_graphviz_dir`<br>`dump_graphviz_video` | *removed (instead, visualize models using [Netron](https://lutzroeder.github.io/netron/) or [visualize.py](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/tools/visualize.py))*<br><br><br> |\n",
        "`output_format`<br>`drop_control_dependency` | *removed (unsupported features in TF2)*<br><br> |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YdZSoIXEbhg-"
      },
      "source": [
        "##  Examples\n",
        "\n",
        "You'll now walk through some examples to convert legacy TF1 models to TF1/TF2 SavedModels and then convert them to TF2 TFLite models.\n",
        "\n",
        "### Setup\n",
        "\n",
        "Start with the necessary TensorFlow imports."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iE0vSfMXumKI"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "import tensorflow.compat.v1 as tf1\n",
        "import numpy as np\n",
        "\n",
        "import logging\n",
        "logger = tf.get_logger()\n",
        "logger.setLevel(logging.ERROR)\n",
        "\n",
        "import shutil\n",
        "def remove_dir(path):\n",
        "  try:\n",
        "    shutil.rmtree(path)\n",
        "  except:\n",
        "    pass"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "89VllCprnFto"
      },
      "source": [
        "Create all the necessary TF1 model formats."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Bwq8EFiwjzjx"
      },
      "outputs": [],
      "source": [
        "# Create a TF1 SavedModel\n",
        "SAVED_MODEL_DIR = \"tf_saved_model/\"\n",
        "remove_dir(SAVED_MODEL_DIR)\n",
        "with tf1.Graph().as_default() as g:\n",
        "  with tf1.Session() as sess:\n",
        "    input = tf1.placeholder(tf.float32, shape=(3,), name='input')\n",
        "    output = input + 2\n",
        "    # print(\"result: \", sess.run(output, {input: [0., 2., 4.]}))\n",
        "    tf1.saved_model.simple_save(\n",
        "        sess, SAVED_MODEL_DIR,\n",
        "        inputs={'input': input}, \n",
        "        outputs={'output': output})\n",
        "print(\"TF1 SavedModel path: \", SAVED_MODEL_DIR)\n",
        "\n",
        "# Create a TF1 Keras model\n",
        "KERAS_MODEL_PATH = 'tf_keras_model.h5'\n",
        "model = tf1.keras.models.Sequential([\n",
        "    tf1.keras.layers.InputLayer(input_shape=(128, 128, 3,), name='input'),\n",
        "    tf1.keras.layers.Dense(units=16, input_shape=(128, 128, 3,), activation='relu'),\n",
        "    tf1.keras.layers.Dense(units=1, name='output')\n",
        "])\n",
        "model.save(KERAS_MODEL_PATH, save_format='h5')\n",
        "print(\"TF1 Keras Model path: \", KERAS_MODEL_PATH)\n",
        "\n",
        "# Create a TF1 frozen GraphDef model\n",
        "GRAPH_DEF_MODEL_PATH = tf.keras.utils.get_file(\n",
        "    'mobilenet_v1_0.25_128',\n",
        "    origin='https://storage.googleapis.com/download.tensorflow.org/models/mobilenet_v1_0.25_128_frozen.tgz',\n",
        "    untar=True,\n",
        ") + '/frozen_graph.pb'\n",
        "\n",
        "print(\"TF1 frozen GraphDef path: \", GRAPH_DEF_MODEL_PATH)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EzMBpG5rdt-7"
      },
      "source": [
        "### 1. Convert a TF1 SavedModel to a TFLite model\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GFWIlVridt_F"
      },
      "source": [
        "#### Before: Converting with TF1\n",
        "This is typical code for TF1-style TFlite conversion.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dzXHHBQRdt_F"
      },
      "outputs": [],
      "source": [
        "converter = tf1.lite.TFLiteConverter.from_saved_model(\n",
        "    saved_model_dir=SAVED_MODEL_DIR,\n",
        "    input_arrays=['input'],\n",
        "    input_shapes={'input' : [3]}\n",
        ")\n",
        "converter.optimizations = {tf.lite.Optimize.DEFAULT}\n",
        "converter.change_concat_input_ranges = True\n",
        "tflite_model = converter.convert()\n",
        "# Ignore warning: \"Use '@tf.function' or '@defun' to decorate the function.\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NUptsxK_MUy2"
      },
      "source": [
        "#### After: Converting with TF2\n",
        "\n",
        "Directly convert the TF1 SavedModel to a TFLite model, with a smaller v2 converter flags set."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0OyBjZ6Kdt_F"
      },
      "outputs": [],
      "source": [
        "# Convert TF1 SavedModel to a TFLite model.\n",
        "converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir=SAVED_MODEL_DIR)\n",
        "converter.optimizations = {tf.lite.Optimize.DEFAULT}\n",
        "tflite_model = converter.convert()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yiwu3sso__fH"
      },
      "source": [
        "### 2. Convert a TF1 Keras model file to a TFLite model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9WTPvPih__fR"
      },
      "source": [
        "#### Before: Converting with TF1\n",
        "This is typical code for TF1-style TFlite conversion."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9EXO0xYq__fR"
      },
      "outputs": [],
      "source": [
        "converter = tf1.lite.TFLiteConverter.from_keras_model_file(model_file=KERAS_MODEL_PATH)\n",
        "converter.optimizations = {tf.lite.Optimize.DEFAULT}\n",
        "converter.change_concat_input_ranges = True\n",
        "tflite_model = converter.convert()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9l6ppTtTZ5Bz"
      },
      "source": [
        "#### After: Converting with TF2\n",
        "\n",
        "First, convert the TF1 Keras model file to a TF2 SavedModel and then convert it to a TFLite model, with a smaller v2 converter flags set."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IGB5ZMGl__fR"
      },
      "outputs": [],
      "source": [
        "# Convert TF1 Keras model file to TF2 SavedModel.\n",
        "model = tf.keras.models.load_model(KERAS_MODEL_PATH)\n",
        "model.save(filepath='saved_model_2/')\n",
        "\n",
        "# Convert TF2 SavedModel to a TFLite model.\n",
        "converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir='saved_model_2/')\n",
        "tflite_model = converter.convert()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "v5Zf6G4M-sZz"
      },
      "source": [
        "### 3. Convert a TF1 frozen GraphDef to a TFLite model\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DzCJOV7AUlGZ"
      },
      "source": [
        "#### Before: Converting with TF1\n",
        "\n",
        "This is typical code for TF1-style TFlite conversion."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "r7RvcdRv6lll"
      },
      "outputs": [],
      "source": [
        "converter = tf1.lite.TFLiteConverter.from_frozen_graph(\n",
        "    graph_def_file=GRAPH_DEF_MODEL_PATH,\n",
        "    input_arrays=['input'],\n",
        "    input_shapes={'input' : [1, 128, 128, 3]},\n",
        "    output_arrays=['MobilenetV1/Predictions/Softmax'],\n",
        ")\n",
        "converter.optimizations = {tf.lite.Optimize.DEFAULT}\n",
        "converter.change_concat_input_ranges = True\n",
        "tflite_model = converter.convert()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZdIogJsKaMNH"
      },
      "source": [
        "#### After: Converting with TF2\n",
        "\n",
        "First, convert the TF1 frozen GraphDef to a TF1 SavedModel and then convert it to a TFLite model, with a smaller v2 converter flags set.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Oigap0TZxjWG"
      },
      "outputs": [],
      "source": [
        "## Convert TF1 frozen Graph to TF1 SavedModel.\n",
        "\n",
        "# Load the graph as a v1.GraphDef\n",
        "import pathlib\n",
        "gdef = tf.compat.v1.GraphDef()\n",
        "gdef.ParseFromString(pathlib.Path(GRAPH_DEF_MODEL_PATH).read_bytes())\n",
        "\n",
        "# Convert the GraphDef to a tf.Graph\n",
        "with tf.Graph().as_default() as g:\n",
        "  tf.graph_util.import_graph_def(gdef, name=\"\")\n",
        "\n",
        "# Look up the input and output tensors.\n",
        "input_tensor = g.get_tensor_by_name('input:0') \n",
        "output_tensor = g.get_tensor_by_name('MobilenetV1/Predictions/Softmax:0')\n",
        "\n",
        "# Save the graph as a TF1 Savedmodel\n",
        "remove_dir('saved_model_3/')\n",
        "with tf.compat.v1.Session(graph=g) as s:\n",
        "  tf.compat.v1.saved_model.simple_save(\n",
        "      session=s,\n",
        "      export_dir='saved_model_3/',\n",
        "      inputs={'input':input_tensor},\n",
        "      outputs={'output':output_tensor})\n",
        "\n",
        "# Convert TF1 SavedModel to a TFLite model.\n",
        "converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir='saved_model_3/')\n",
        "converter.optimizations = {tf.lite.Optimize.DEFAULT}\n",
        "tflite_model = converter.convert()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MFbsddkOw4Wl"
      },
      "source": [
        "# Further reading\n",
        "\n",
        "*   Refer to the [TFLite Guide](https://www.tensorflow.org/lite/guide) to learn more about the workflows and latest features.\n",
        "* If you're using  TF1 code or legacy TF1 model formats (Keras `.h5` files, frozen GraphDef `.pb`, etc), please update your code and migrate your models to the [TF2 SavedModel format](https://www.tensorflow.org/guide/saved_model). \n"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "tflite.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
