retrain_classification_qat_tf1.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "license"
      },
      "source": [
        "##### *Copyright 2020 Google LLC*\n",
        "*Licensed under the Apache License, Version 2.0 (the \"License\")*"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "both",
        "id": "rKwqeqWBXANA"
      },
      "outputs": [],
      "source": [
        "# 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": "hRTa3Ee15WsJ"
      },
      "source": [
        "# Retrain a classification model for Edge TPU with quant-aware training (TF 1.15)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TaX0smDP7xQY"
      },
      "source": [
        "This notebook uses a set of TensorFlow training scripts to perform transfer-learning on a quantization-aware classification model and then convert it for compatibility with the [Edge TPU](https://coral.ai/products/).\n",
        "\n",
        "Specifically, this tutorial shows you how to perform [fine-tuning](https://github.com/tensorflow/models/blob/master/research/slim/README.md#fine-tuning-a-model-from-an-existing-checkpoint) on the MobileNet V1 model so it can recognize a new set of classes (five types of flowers), using TensorFlow r1.15.\n",
        "\n",
        "Beware that, compared to a desktop computer, this training can take much longer in Colab because Colab provides limited resources for long-running operations. So you'll likely see faster training speeds if you [connect this notebook to a local runtime](https://research.google.com/colaboratory/local-runtimes.html), or instead follow the [tutorial to run this training in Docker](https://coral.ai/docs/edgetpu/retrain-classification/) (which includes more documentation about this process).\n",
        "\n",
        "**Note:** As an alternative approach, check out this other Colab notebook that [retrains a classification model for the Edge TPU using post-training quantization](https://colab.sandbox.google.com/github/google-coral/tutorials/blob/master/retrain_classification_ptq_tf1.ipynb#scrollTo=hRTa3Ee15WsJ) (instead of quantization-aware training)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GTCYQg_be8C0"
      },
      "source": [
        "## Import TensorFlow"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kxAceTA36NKQ"
      },
      "outputs": [],
      "source": [
        "! pip uninstall tensorflow -y\n",
        "! pip install tensorflow-gpu==1.15"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5Ebgg83X9oTh"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "print(tf.__version__)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gpTmoIxuranU"
      },
      "source": [
        "## Clone the model and training repos"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "r_zobAPP8J9Y"
      },
      "outputs": [],
      "source": [
        "! git clone https://github.com/tensorflow/models.git"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8dagXBMpej0E"
      },
      "outputs": [],
      "source": [
        "! cd models \u0026\u0026 git checkout f788046ca876a8820e05b0b48c1fc2e16b0955bc"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Vy4Q_Uva9eii"
      },
      "outputs": [],
      "source": [
        "! git clone https://github.com/google-coral/tutorials.git\n",
        "\n",
        "! cp tutorials/docker/classification/scripts/* models/research/slim/"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IweNl64rridS"
      },
      "source": [
        "## Convert training data to TFRecord"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IBJgQgqmAFKU"
      },
      "outputs": [],
      "source": [
        "%cd models/research/slim"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Mbz5nKlDAorQ"
      },
      "outputs": [],
      "source": [
        "! ./prepare_checkpoint_and_dataset.sh --network_type mobilenet_v1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kg3oMLs1rus7"
      },
      "source": [
        "## Perform transfer-learning"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qePNPDZQGNC6"
      },
      "source": [
        "The following script takes a couple hours to finish in Colab. (You can shorten by reducing the steps, but that reduces the final accuracy.)\n",
        "\n",
        "If you didn't already select \"Run all\" then you should run all remaining cells now. That will ensure the rest of the notebook completes while you are away, avoiding the chance that the Colab runtime times-out and you lose the training data before you download the model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "au3At-_QGZey"
      },
      "outputs": [],
      "source": [
        "%env NUM_TRAINING_STEPS=300"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ISp0a9D7Ap4A"
      },
      "outputs": [],
      "source": [
        "! ./start_training.sh --network_type mobilenet_v1 --train_steps $NUM_TRAINING_STEPS"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "11y-jGfZLbNB"
      },
      "source": [
        "As training progresses, you can see new checkpoint files appear in the `models/research/slim/transfer_learn/train/` directory."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7Y0CmosGAywn"
      },
      "outputs": [],
      "source": [
        "! ./run_evaluation.sh --network_type mobilenet_v1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Quv4hQWNhaAH"
      },
      "source": [
        "## Compile for the Edge TPU"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jq9z4ctFiwp6"
      },
      "outputs": [],
      "source": [
        "! ./convert_checkpoint_to_edgetpu_tflite.sh --network_type mobilenet_v1 --checkpoint_num $NUM_TRAINING_STEPS"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6RxtslKJf2td"
      },
      "outputs": [],
      "source": [
        "! curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key add -\n",
        "\n",
        "! echo \"deb https://packages.cloud.google.com/apt coral-edgetpu-stable main\" | sudo tee /etc/apt/sources.list.d/coral-edgetpu.list\n",
        "\n",
        "! sudo apt-get update\n",
        "\n",
        "! sudo apt-get install edgetpu-compiler\t"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9doQdA2QkPnV"
      },
      "outputs": [],
      "source": [
        "%cd transfer_learn/models/\n",
        "\n",
        "! ls"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TYUhhhMdkbrY"
      },
      "outputs": [],
      "source": [
        "! edgetpu_compiler output_tflite_graph.tflite"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NPt8RdopXsZv"
      },
      "source": [
        "Download the files:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DtuE-CnPkdfI"
      },
      "outputs": [],
      "source": [
        "from google.colab import files\n",
        "\n",
        "files.download('output_tflite_graph_edgetpu.tflite')\n",
        "files.download('labels.txt')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_qOCP3mXXvsm"
      },
      "source": [
        "If you get a \"Failed to fetch\" error here, it's probably because the files weren't done saving. So just wait a moment and try again.\n",
        "\n",
        "Also look out for a browser popup that might need approval to download the files."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_TZTwG7nhm0C"
      },
      "source": [
        "## Run the model on the Edge TPU\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RwywT4ZpQjLf"
      },
      "source": [
        "You can now run the model on your Coral device with acceleration on the Edge TPU.\n",
        "\n",
        "First, you need some flower photos to try. Remember that you've trained this model to recognize just five flower classes: daisy,\n",
        "dandelion, roses, sunflowers, and tulips. Here are a couple images you can try (provided by the\n",
        "[Open Images Dataset](https://storage.googleapis.com/openimages/web/index.html)):\n",
        "\n",
        "```\n",
        "wget https://c4.staticflickr.com/3/2856/13169252123_e4c5086ea3_z.jpg -O flower.jpg \u0026\u0026 \\\n",
        "wget https://c6.staticflickr.com/4/3372/3416475881_726f0d33fe_z.jpg -O flower2.jpg\n",
        "```\n",
        "\n",
        "Then, try running an inference using [this example code for the PyCoral API](https://github.com/google-coral/pycoral/blob/master/examples/classify_image.py). Just clone that repo and run the script using the model files you downloaded above (also be sure you have [installed the PyCoral API](https://coral.ai/software/#pycoral-api)):\n",
        "\n",
        "```\n",
        "git clone https://github.com/google-coral/pycoral\n",
        "\n",
        "cd pycoral/examples/\n",
        "\n",
        "python3 classify_image.py \\\n",
        "  --model output_tflite_graph_edgetpu.tflite  \\\n",
        "  --labels labels.txt \\\n",
        "  --input flower.jpg\n",
        "```\n",
        "\n",
        "Check out more examples for running inference at [coral.ai/examples](https://coral.ai/examples/#code-examples/)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "p2tyWn83VOAF"
      },
      "source": [
        "## Implementation details\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "x5tKMtKVVDps"
      },
      "source": [
        "\n",
        "All the scripts used in this notebook come from the following locations:\u003cbr\u003e\n",
        "+  https://github.com/google-coral/tutorials/tree/master/docker/classification/scripts\n",
        "+  https://github.com/tensorflow/models/tree/r1.13.0/research/slim\n",
        "\n",
        "More explanation of the steps in this tutorial is available at https://coral.ai/docs/edgetpu/retrain-classification/."
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "license"
      ],
      "name": "Retrain MobileNet V1 classifier for the Edge TPU (TF1, quant-aware)",
      "private_outputs": true,
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}