retrain_detection_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 detection 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 object detection model and then convert it for compatibility with the [Edge TPU](https://coral.ai/products/).\n",
        "\n",
        "Specifically, this tutorial shows you how to retrain a MobileNet V1 SSD model so that it detects two pets: Abyssinian cats and American Bulldogs (from the [Oxford-IIIT Pets Dataset](https://www.robots.ox.ac.uk/~vgg/data/pets/)), using TensorFlow r1.15.\n",
        "\n",
        "Beware that, compared to a desktop computer, this training can take *a lot* 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-detection/) (which includes more documentation about this process)."
      ]
    },
    {
      "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\n",
        "! pip install -U numpy==1.19.5\n",
        "! pip install -U pycocotools==2.0.1"
      ]
    },
    {
      "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": "U4Yftz8HsilF"
      },
      "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 -r tutorials/docker/object_detection/scripts/* models/research/"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0Iv-kpe2Xe69"
      },
      "source": [
        "## Import dependencies"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ImucOu0qgMv_"
      },
      "source": [
        "For details, see https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/installation.md"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LJEhkUxlfhh4"
      },
      "outputs": [],
      "source": [
        "! apt-get install -y python python-tk\n",
        "! pip install Cython contextlib2 pillow lxml jupyter matplotlib"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "45oRT7h6XhgP"
      },
      "outputs": [],
      "source": [
        "# Get protoc 3.0.0, rather than the old version already in the container\n",
        "! wget https://www.github.com/google/protobuf/releases/download/v3.0.0/protoc-3.0.0-linux-x86_64.zip\n",
        "! unzip protoc-3.0.0-linux-x86_64.zip -d proto3\n",
        "! mkdir -p local/bin \u0026\u0026 mkdir -p local/include\n",
        "! mv proto3/bin/* local/bin\n",
        "! mv proto3/include/* local/include\n",
        "! rm -rf proto3 protoc-3.0.0-linux-x86_64.zip"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "snUGUfruaTFa"
      },
      "outputs": [],
      "source": [
        "# Install pycocoapi\n",
        "! git clone --depth 1 https://github.com/cocodataset/cocoapi.git\n",
        "! (cd cocoapi/PythonAPI \u0026\u0026 make -j8)\n",
        "! cp -r cocoapi/PythonAPI/pycocotools/ models/research/\n",
        "! rm -rf cocoapi"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Yz0nN7eVeXo6"
      },
      "outputs": [],
      "source": [
        "# Run protoc on the object detection repo (generate .py files from .proto)\n",
        "% cd models/research/\n",
        "! ../../local/bin/protoc object_detection/protos/*.proto --python_out=."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "53D-U0_gg8VB"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "os.environ['PYTHONPATH'] += \":/content/models/research:/content/models/research/slim\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PpXtNIFxkms2"
      },
      "source": [
        "Just to verify everything is correctly set up:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ftH0svNxgUm4"
      },
      "outputs": [],
      "source": [
        "! python object_detection/builders/model_builder_test.py"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IweNl64rridS"
      },
      "source": [
        "## Convert training data to TFRecord"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IdGWra9PBA-9"
      },
      "source": [
        "To train with different images, read [how to configure your own training data](https://coral.ai/docs/edgetpu/retrain-detection/#configure-your-own-training-data)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Mbz5nKlDAorQ"
      },
      "outputs": [],
      "source": [
        "! ./prepare_checkpoint_and_dataset.sh --network_type mobilenet_v1_ssd --train_whole_model false"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kg3oMLs1rus7"
      },
      "source": [
        "## Perform transfer-learning"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "M2FIkwyhW8IX"
      },
      "source": [
        "The following script takes several 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": "8NkqCq8g9A5M"
      },
      "outputs": [],
      "source": [
        "%env NUM_TRAINING_STEPS=500\n",
        "%env NUM_EVAL_STEPS=100\n",
        "\n",
        "# If you're retraining the whole model, we suggest thes values:\n",
        "# %env NUM_TRAINING_STEPS=50000\n",
        "# %env NUM_EVAL_STEPS=2000"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ISp0a9D7Ap4A"
      },
      "outputs": [],
      "source": [
        "! ./retrain_detection_model.sh --num_training_steps $NUM_TRAINING_STEPS --num_eval_steps $NUM_EVAL_STEPS"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "G1jjZIKpmbyB"
      },
      "source": [
        "As training progresses, you can see new checkpoint files appear in the `models/research/learn_pet/train/` directory."
      ]
    },
    {
      "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 --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 learn_pet/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, find some photos to try. Remember that you've trained this model to recognize just two classes: Abyssinian cats and\n",
        "American Bulldogs. So here are a couple images that should provide results (provided by the\n",
        "[Open Images Dataset](https://storage.googleapis.com/openimages/web/index.html)): \n",
        "\n",
        "```\n",
        "wget https://c4.staticflickr.com/8/7580/15865399370_ffa5b49d20_z.jpg -O dog.jpg \u0026\u0026 \\\n",
        "wget https://c6.staticflickr.com/9/8534/8652503705_687d957a29_z.jpg -O cat.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/detect_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 detect_image.py \\\n",
        "  --model output_tflite_graph_edgetpu.tflite \\\n",
        "  --labels labels.txt \\\n",
        "  --input dog.jpg \\\n",
        "  --output dog_result.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/object_detection/scripts\n",
        "+  https://github.com/tensorflow/models/tree/r1.13.0/research/object_detection/\n",
        "\n",
        "More explanation of the steps in this tutorial is available at\n",
        "https://coral.ai/docs/edgetpu/retrain-detection/."
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "license"
      ],
      "name": "Retrain SSD MobileNet V1 detector for the Edge TPU (TF1)",
      "private_outputs": true,
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}