30420902-Retraining-an-image-classifier-template.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "Retraining_an_image_classifier_template.ipynb",
      "provenance": [],
      "private_outputs": true,
      "collapsed_sections": [
        "ScitaPqhKtuW"
      ],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ScitaPqhKtuW"
      },
      "source": [
        ""
      ]
    },
    {
      "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"
      ]
    },
    {
      "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",
        "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",
        "cellView": "form"
      },
      "source": [
        "module_selection = (\"inception_v3\", 299) #@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 dataset\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pRUe01fg1wmG"
      },
      "source": [
        "zip_file_name = \"\" #@param {type:\"string\"}\n",
        "!unzip {zip_file_name}.zip"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8h1bxmGDqjvq"
      },
      "source": [
        "data_dir = '/content/'+zip_file_name"
      ],
      "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",
        "zip_file_name = \"\"\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": "QzW4oNRjILaq"
      },
      "source": [
        "## Deployment to TensorFlow Lite\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CuZuFJJOnIHF"
      },
      "source": [
        "**convert the model to tflite format**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oPve3m5xlfdX"
      },
      "source": [
        "converter = tf.lite.TFLiteConverter.from_keras_model(model)\n",
        "tflite_model = converter.convert()\n",
        "open(\"model.tflite\",\"wb\").write(tflite_model)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nO5yn4xH7QRL"
      },
      "source": [
        "labels_list = []\n",
        "for class_string, class_index in valid_generator.class_indices.items():\n",
        "      print(class_string)\n",
        "      labels_list.append(class_string)\n",
        "with open('labels.txt', 'w') as f:\n",
        "    for item in labels_list:\n",
        "        f.write(\"%s\\n\" % item)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dBKHH2S3x7ku"
      },
      "source": [
        "**Evaluate the model**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_wqEmD0xIqeG"
      },
      "source": [
        "# docs_infra: no_execute\n",
        "interpreter = tf.lite.Interpreter(model_content=tflite_model)\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",
        "cellView": "form"
      },
      "source": [
        "#@markdown For rapid experimentation, start with a moderate number of examples.\n",
        "# docs_infra: no_execute\n",
        "num_eval_examples = 125  #@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": []
    }
  ]
}