TF2_Retraining_an_image_classifier.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "Copy of TF Hub for TF2: Retraining an image classifier",
      "private_outputs": true,
      "provenance": [],
      "collapsed_sections": [
        "ScitaPqhKtuW"
      ],
      "toc_visible": true,
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/sxl613/Machine-Learning/blob/master/TF2_Retraining_an_image_classifier.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ScitaPqhKtuW"
      },
      "source": [
        "##### Copyright 2019 The TensorFlow Hub Authors.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jvztxQ6VsK2k"
      },
      "source": [
        "# Copyright 2019 The TensorFlow Hub Authors. All Rights Reserved.\n",
        "#\n",
        "# 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",
        "#     http://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.\n",
        "# =============================================================================="
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oYM61xrTsP5d"
      },
      "source": [
        "# Retraining an Image Classifier\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MfBg1C5NB3X0"
      },
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://www.tensorflow.org/hub/tutorials/tf2_image_retraining\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />View on TensorFlow.org</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/hub/blob/master/examples/colab/tf2_image_retraining.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/tensorflow/hub/blob/master/examples/colab/tf2_image_retraining.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View on GitHub</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://storage.googleapis.com/tensorflow_docs/hub/examples/colab/tf2_image_retraining.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\" />Download notebook</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://tfhub.dev/google/imagenet/mobilenet_v2_100_224/feature_vector/4\"><img src=\"https://www.tensorflow.org/images/hub_logo_32px.png\" />See TF Hub model</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L1otmJgmbahf"
      },
      "source": [
        "## Introduction\n",
        "\n",
        "Image classification models have millions of parameters. Training them from\n",
        "scratch requires a lot of labeled training data and a lot of computing power. Transfer learning is a technique that shortcuts much of this by taking a piece of a model that has already been trained on a related task and reusing it in a new model.\n",
        "\n",
        "This Colab demonstrates how to build a Keras model for classifying five species of flowers by using a pre-trained TF2 SavedModel from TensorFlow Hub for image feature extraction, trained on the much larger and more general ImageNet dataset. Optionally, the feature extractor can be trained (\"fine-tuned\") alongside the newly added classifier.\n",
        "\n",
        "### Looking for a tool instead?\n",
        "\n",
        "This is a TensorFlow coding tutorial. If you want a tool that just builds the TensorFlow or TF Lite model for, take a look at the [make_image_classifier](https://github.com/tensorflow/hub/tree/master/tensorflow_hub/tools/make_image_classifier) command-line tool that gets [installed](https://www.tensorflow.org/hub/installation) by the PIP package `tensorflow-hub[make_image_classifier]`, or at [this](https://colab.sandbox.google.com/github/tensorflow/examples/blob/master/tensorflow_examples/lite/model_maker/demo/image_classification.ipynb) TF Lite colab.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bL54LWCHt5q5"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dlauq-4FWGZM"
      },
      "source": [
        "import itertools\n",
        "import os\n",
        "\n",
        "import matplotlib.pylab as plt\n",
        "import numpy as np\n",
        "\n",
        "import tensorflow as tf\n",
        "import tensorflow_hub as hub\n",
        "\n",
        "print(\"TF version:\", tf.__version__)\n",
        "print(\"Hub version:\", hub.__version__)\n",
        "print(\"GPU is\", \"available\" if tf.test.is_gpu_available() else \"NOT AVAILABLE\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mmaHHH7Pvmth"
      },
      "source": [
        "## Select the TF2 SavedModel module to use\n",
        "\n",
        "For starters, use [https://tfhub.dev/google/imagenet/mobilenet_v2_100_224/feature_vector/4](https://tfhub.dev/google/imagenet/mobilenet_v2_100_224/feature_vector/4). The same URL can be used in code to identify the SavedModel and in your browser to show its documentation. (Note that models in TF1 Hub format won't work here.)\n",
        "\n",
        "You can find more TF2 models that generate image feature vectors [here](https://tfhub.dev/s?module-type=image-feature-vector&tf-version=tf2)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FlsEcKVeuCnf"
      },
      "source": [
        "module_selection = (\"mobilenet_v2_100_224\", 224) #@param [\"(\\\"mobilenet_v2_100_224\\\", 224)\", \"(\\\"inception_v3\\\", 299)\"] {type:\"raw\", allow-input: true}\n",
        "handle_base, pixels = module_selection\n",
        "MODULE_HANDLE =\"https://tfhub.dev/google/imagenet/{}/feature_vector/4\".format(handle_base)\n",
        "IMAGE_SIZE = (pixels, pixels)\n",
        "print(\"Using {} with input size {}\".format(MODULE_HANDLE, IMAGE_SIZE))\n",
        "\n",
        "BATCH_SIZE = 32 #@param {type:\"integer\"}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yTY8qzyYv3vl"
      },
      "source": [
        "## Set up the Flowers dataset\n",
        "\n",
        "Inputs are suitably resized for the selected module. Dataset augmentation (i.e., random distortions of an image each time it is read) improves training, esp. when fine-tuning."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WBtFK1hO8KsO"
      },
      "source": [
        "data_dir = tf.keras.utils.get_file(\n",
        "    'flower_photos',\n",
        "    'https://storage.googleapis.com/download.tensorflow.org/example_images/flower_photos.tgz',\n",
        "    untar=True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "umB5tswsfTEQ"
      },
      "source": [
        "datagen_kwargs = dict(rescale=1./255, validation_split=.20)\n",
        "dataflow_kwargs = dict(target_size=IMAGE_SIZE, batch_size=BATCH_SIZE,\n",
        "                   interpolation=\"bilinear\")\n",
        "\n",
        "valid_datagen = tf.keras.preprocessing.image.ImageDataGenerator(\n",
        "    **datagen_kwargs)\n",
        "valid_generator = valid_datagen.flow_from_directory(\n",
        "    data_dir, subset=\"validation\", shuffle=False, **dataflow_kwargs)\n",
        "\n",
        "do_data_augmentation = False #@param {type:\"boolean\"}\n",
        "if do_data_augmentation:\n",
        "  train_datagen = tf.keras.preprocessing.image.ImageDataGenerator(\n",
        "      rotation_range=40,\n",
        "      horizontal_flip=True,\n",
        "      width_shift_range=0.2, height_shift_range=0.2,\n",
        "      shear_range=0.2, zoom_range=0.2,\n",
        "      **datagen_kwargs)\n",
        "else:\n",
        "  train_datagen = valid_datagen\n",
        "train_generator = train_datagen.flow_from_directory(\n",
        "    data_dir, subset=\"training\", shuffle=True, **dataflow_kwargs)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FS_gVStowW3G"
      },
      "source": [
        "## Defining the model\n",
        "\n",
        "All it takes is to put a linear classifier on top of the `feature_extractor_layer` with the Hub module.\n",
        "\n",
        "For speed, we start out with a non-trainable `feature_extractor_layer`, but you can also enable fine-tuning for greater accuracy."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RaJW3XrPyFiF"
      },
      "source": [
        "do_fine_tuning = False #@param {type:\"boolean\"}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "50FYNIb1dmJH"
      },
      "source": [
        "print(\"Building model with\", MODULE_HANDLE)\n",
        "model = tf.keras.Sequential([\n",
        "    # Explicitly define the input shape so the model can be properly\n",
        "    # loaded by the TFLiteConverter\n",
        "    tf.keras.layers.InputLayer(input_shape=IMAGE_SIZE + (3,)),\n",
        "    hub.KerasLayer(MODULE_HANDLE, trainable=do_fine_tuning),\n",
        "    tf.keras.layers.Dropout(rate=0.2),\n",
        "    tf.keras.layers.Dense(train_generator.num_classes,\n",
        "                          kernel_regularizer=tf.keras.regularizers.l2(0.0001))\n",
        "])\n",
        "model.build((None,)+IMAGE_SIZE+(3,))\n",
        "model.summary()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "u2e5WupIw2N2"
      },
      "source": [
        "## Training the model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9f3yBUvkd_VJ"
      },
      "source": [
        "model.compile(\n",
        "  optimizer=tf.keras.optimizers.SGD(lr=0.005, momentum=0.9), \n",
        "  loss=tf.keras.losses.CategoricalCrossentropy(from_logits=True, label_smoothing=0.1),\n",
        "  metrics=['accuracy'])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "w_YKX2Qnfg6x"
      },
      "source": [
        "steps_per_epoch = train_generator.samples // train_generator.batch_size\n",
        "validation_steps = valid_generator.samples // valid_generator.batch_size\n",
        "hist = model.fit(\n",
        "    train_generator,\n",
        "    epochs=5, steps_per_epoch=steps_per_epoch,\n",
        "    validation_data=valid_generator,\n",
        "    validation_steps=validation_steps).history"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CYOw0fTO1W4x"
      },
      "source": [
        "plt.figure()\n",
        "plt.ylabel(\"Loss (training and validation)\")\n",
        "plt.xlabel(\"Training Steps\")\n",
        "plt.ylim([0,2])\n",
        "plt.plot(hist[\"loss\"])\n",
        "plt.plot(hist[\"val_loss\"])\n",
        "\n",
        "plt.figure()\n",
        "plt.ylabel(\"Accuracy (training and validation)\")\n",
        "plt.xlabel(\"Training Steps\")\n",
        "plt.ylim([0,1])\n",
        "plt.plot(hist[\"accuracy\"])\n",
        "plt.plot(hist[\"val_accuracy\"])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jZ8DKKgeKv4-"
      },
      "source": [
        "Try out the model on an image from the validation data:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oi1iCNB9K1Ai"
      },
      "source": [
        "def get_class_string_from_index(index):\n",
        "   for class_string, class_index in valid_generator.class_indices.items():\n",
        "      if class_index == index:\n",
        "         return class_string\n",
        "\n",
        "x, y = next(valid_generator)\n",
        "image = x[0, :, :, :]\n",
        "true_index = np.argmax(y[0])\n",
        "plt.imshow(image)\n",
        "plt.axis('off')\n",
        "plt.show()\n",
        "\n",
        "# Expand the validation image to (1, 224, 224, 3) before predicting the label\n",
        "prediction_scores = model.predict(np.expand_dims(image, axis=0))\n",
        "predicted_index = np.argmax(prediction_scores)\n",
        "print(\"True label: \" + get_class_string_from_index(true_index))\n",
        "print(\"Predicted label: \" + get_class_string_from_index(predicted_index))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YCsAsQM1IRvA"
      },
      "source": [
        "Finally, the trained model can be saved for deployment to TF Serving or TF Lite (on mobile) as follows."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LGvTi69oIc2d"
      },
      "source": [
        "saved_model_path = \"/tmp/saved_flowers_model\"\n",
        "tf.saved_model.save(model, saved_model_path)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QzW4oNRjILaq"
      },
      "source": [
        "## Optional: Deployment to TensorFlow Lite\n",
        "\n",
        "[TensorFlow Lite](https://www.tensorflow.org/lite) lets you deploy TensorFlow models to mobile and IoT devices. The code below shows how to convert the trained model to TF Lite and apply post-training tools from the [TensorFlow Model Optimization Toolkit](https://www.tensorflow.org/model_optimization). Finally, it runs it in the TF Lite Interpreter to examine the resulting quality\n",
        "\n",
        "  * Converting without optimization provides the same results as before (up to roundoff error).\n",
        "  * Converting with optimization without any data quantizes the model weights to 8 bits, but inference still uses floating-point computation for the neural network activations. This reduces model size almost by a factor of 4 and improves CPU latency on mobile devices.\n",
        "  * On top, computation of the neural network activations can be quantized to 8-bit integers as well if a small reference dataset is provided to calibrate the quantization range. On a mobile device, this accelerates inference further and makes it possible to run on accelerators like EdgeTPU."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Va1Vo92fSyV6"
      },
      "source": [
        "#@title Optimization settings\n",
        "# docs_infra: no_execute\n",
        "# TODO(b/156102192)\n",
        "optimize_lite_model = False  #@param {type:\"boolean\"}\n",
        "#@markdown Setting a value greater than zero enables quantization of neural network activations. A few dozen is already a useful amount.\n",
        "num_calibration_examples = 60  #@param {type:\"slider\", min:0, max:1000, step:1}\n",
        "representative_dataset = None\n",
        "if optimize_lite_model and num_calibration_examples:\n",
        "  # Use a bounded number of training examples without labels for calibration.\n",
        "  # TFLiteConverter expects a list of input tensors, each with batch size 1.\n",
        "  representative_dataset = lambda: itertools.islice(\n",
        "      ([image[None, ...]] for batch, _ in train_generator for image in batch),\n",
        "      num_calibration_examples)\n",
        "\n",
        "converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_path)\n",
        "if optimize_lite_model:\n",
        "  converter.optimizations = [tf.lite.Optimize.DEFAULT]\n",
        "  if representative_dataset:  # This is optional, see above.\n",
        "    converter.representative_dataset = representative_dataset\n",
        "lite_model_content = converter.convert()\n",
        "\n",
        "with open(\"/tmp/lite_flowers_model\", \"wb\") as f:\n",
        "  f.write(lite_model_content)\n",
        "print(\"Wrote %sTFLite model of %d bytes.\" %\n",
        "      (\"optimized \" if optimize_lite_model else \"\", len(lite_model_content)))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_wqEmD0xIqeG"
      },
      "source": [
        "# docs_infra: no_execute\n",
        "interpreter = tf.lite.Interpreter(model_content=lite_model_content)\n",
        "# This little helper wraps the TF Lite interpreter as a numpy-to-numpy function.\n",
        "def lite_model(images):\n",
        "  interpreter.allocate_tensors()\n",
        "  interpreter.set_tensor(interpreter.get_input_details()[0]['index'], images)\n",
        "  interpreter.invoke()\n",
        "  return interpreter.get_tensor(interpreter.get_output_details()[0]['index'])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JMMK-fZrKrk8"
      },
      "source": [
        "#@markdown For rapid experimentation, start with a moderate number of examples.\n",
        "# docs_infra: no_execute\n",
        "num_eval_examples = 50  #@param {type:\"slider\", min:0, max:700}\n",
        "eval_dataset = ((image, label)  # TFLite expects batch size 1.\n",
        "                for batch in train_generator\n",
        "                for (image, label) in zip(*batch))\n",
        "count = 0\n",
        "count_lite_tf_agree = 0\n",
        "count_lite_correct = 0\n",
        "for image, label in eval_dataset:\n",
        "  probs_lite = lite_model(image[None, ...])[0]\n",
        "  probs_tf = model(image[None, ...]).numpy()[0]\n",
        "  y_lite = np.argmax(probs_lite)\n",
        "  y_tf = np.argmax(probs_tf)\n",
        "  y_true = np.argmax(label)\n",
        "  count +=1\n",
        "  if y_lite == y_tf: count_lite_tf_agree += 1\n",
        "  if y_lite == y_true: count_lite_correct += 1\n",
        "  if count >= num_eval_examples: break\n",
        "print(\"TF Lite model agrees with original model on %d of %d examples (%g%%).\" %\n",
        "      (count_lite_tf_agree, count, 100.0 * count_lite_tf_agree / count))\n",
        "print(\"TF Lite model is accurate on %d of %d examples (%g%%).\" %\n",
        "      (count_lite_correct, count, 100.0 * count_lite_correct / count))"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}