TPUs_in_Colab.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/iypc-team/CoLab/blob/master/TPUs_in_Colab.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import shutil\n",
        "try: shutil.rmtree('/content/sample_data')\n",
        "except: pass\n",
        "print('Updated: 11/08/2022')"
      ],
      "metadata": {
        "id": "cbHEyAiDb0Mh"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FpvUOuC3j27n"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "print(\"Tensorflow version \" + tf.__version__)\n",
        "\n",
        "try:\n",
        "  tpu = tf.distribute.cluster_resolver.TPUClusterResolver()  # TPU detection\n",
        "  print('Running on TPU ', tpu.cluster_spec().as_dict()['worker'])\n",
        "except ValueError:\n",
        "  raise BaseException('ERROR: Not connected to a TPU runtime; please see the previous cell in this notebook for instructions!')\n",
        "\n",
        "tf.config.experimental_connect_to_cluster(tpu)\n",
        "tf.distribute.TPUStrategy(tpu)\n",
        "tpu_strategy = tf.distribute.TPUStrategy(tpu)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vFIMfPmgQa0h"
      },
      "outputs": [],
      "source": [
        "import re\n",
        "import numpy as np\n",
        "from matplotlib import pyplot as plt"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kvPXiovhi3ZZ"
      },
      "source": [
        "\n",
        "## Input data\n",
        "\n",
        "Our input data is stored on Google Cloud Storage. To more fully use the parallelism TPUs offer us, and to avoid bottlenecking on data transfer, we've stored our input data in TFRecord files, 230 images per file.\n",
        "\n",
        "Below, we make heavy use of `tf.data.experimental.AUTOTUNE` to optimize different parts of input loading.\n",
        "\n",
        "All of these techniques are a bit overkill for our (small) dataset, but demonstrate best practices for using TPUs.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LtAVr-4CP1rp"
      },
      "outputs": [],
      "source": [
        "AUTO = tf.data.experimental.AUTOTUNE\n",
        "\n",
        "IMAGE_SIZE = [331, 331]\n",
        "\n",
        "batch_size = 16 * tpu_strategy.num_replicas_in_sync\n",
        "\n",
        "gcs_pattern = 'gs://flowers-public/tfrecords-jpeg-331x331/*.tfrec'\n",
        "validation_split = 0.19\n",
        "filenames = tf.io.gfile.glob(gcs_pattern)\n",
        "split = len(filenames) - int(len(filenames) * validation_split)\n",
        "train_fns = filenames[:split]\n",
        "validation_fns = filenames[split:]\n",
        "        \n",
        "def parse_tfrecord(example):\n",
        "  features = {\n",
        "    \"image\": tf.io.FixedLenFeature([], tf.string), # tf.string means bytestring\n",
        "    \"class\": tf.io.FixedLenFeature([], tf.int64),  # shape [] means scalar\n",
        "    \"one_hot_class\": tf.io.VarLenFeature(tf.float32),\n",
        "  }\n",
        "  example = tf.io.parse_single_example(example, features)\n",
        "  decoded = tf.image.decode_jpeg(example['image'], channels=3)\n",
        "  normalized = tf.cast(decoded, tf.float32) / 255.0 # convert each 0-255 value to floats in [0, 1] range\n",
        "  image_tensor = tf.reshape(normalized, [*IMAGE_SIZE, 3])\n",
        "  one_hot_class = tf.reshape(tf.sparse.to_dense(example['one_hot_class']), [5])\n",
        "  return image_tensor, one_hot_class\n",
        "\n",
        "def load_dataset(filenames):\n",
        "  # Read from TFRecords. For optimal performance, we interleave reads from multiple files.\n",
        "  records = tf.data.TFRecordDataset(filenames, num_parallel_reads=AUTO)\n",
        "  return records.map(parse_tfrecord, num_parallel_calls=AUTO)\n",
        "\n",
        "def get_training_dataset():\n",
        "  dataset = load_dataset(train_fns)\n",
        "\n",
        "  # Create some additional training images by randomly flipping and\n",
        "  # increasing/decreasing the saturation of images in the training set. \n",
        "  def data_augment(image, one_hot_class):\n",
        "    modified = tf.image.random_flip_left_right(image)\n",
        "    modified = tf.image.random_saturation(modified, 0, 2)\n",
        "    return modified, one_hot_class\n",
        "  augmented = dataset.map(data_augment, num_parallel_calls=AUTO)\n",
        "\n",
        "  # Prefetch the next batch while training (autotune prefetch buffer size).\n",
        "  return augmented.repeat().shuffle(2048).batch(batch_size).prefetch(AUTO) \n",
        "\n",
        "training_dataset = get_training_dataset()\n",
        "validation_dataset = load_dataset(validation_fns).batch(batch_size).prefetch(AUTO)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KTukaIGil7_m"
      },
      "source": [
        "Let's take a peek at the training dataset we've created:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-DwfsJq0l_DJ"
      },
      "outputs": [],
      "source": [
        "CLASSES = ['daisy', 'dandelion', 'roses', 'sunflowers', 'tulips']\n",
        "\n",
        "def display_one_flower(image, title, subplot, color):\n",
        "  plt.subplot(subplot)\n",
        "  plt.axis('off')\n",
        "  plt.imshow(image)\n",
        "  plt.title(title, fontsize=16, color=color)\n",
        "  \n",
        "# If model is provided, use it to generate predictions.\n",
        "def display_nine_flowers(images, titles, title_colors=None):\n",
        "  subplot = 331\n",
        "  plt.figure(figsize=(13,13))\n",
        "  for i in range(9):\n",
        "    color = 'black' if title_colors is None else title_colors[i]\n",
        "    display_one_flower(images[i], titles[i], 331+i, color)\n",
        "  plt.tight_layout()\n",
        "  plt.subplots_adjust(wspace=0.1, hspace=0.1)\n",
        "  plt.show()\n",
        "\n",
        "def get_dataset_iterator(dataset, n_examples):\n",
        "  return dataset.unbatch().batch(n_examples).as_numpy_iterator()\n",
        "\n",
        "training_viz_iterator = get_dataset_iterator(training_dataset, 9)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JTnyd7qfbYr4"
      },
      "outputs": [],
      "source": [
        "# Re-run this cell to show a new batch of images\n",
        "images, classes = next(training_viz_iterator)\n",
        "class_idxs = np.argmax(classes, axis=-1) # transform from one-hot array to class number\n",
        "labels = [CLASSES[idx] for idx in class_idxs]\n",
        "display_nine_flowers(images, labels)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ALtRUlxhw8Vt"
      },
      "source": [
        "## Model\n",
        "To get maxmimum accuracy, we leverage a pretrained image recognition model (here, [Xception](http://openaccess.thecvf.com/content_cvpr_2017/papers/Chollet_Xception_Deep_Learning_CVPR_2017_paper.pdf)). We drop the ImageNet-specific top layers (`include_top=false`), and add a max pooling and a softmax layer to predict our 5 classes."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hJl3vNtJOB-x"
      },
      "outputs": [],
      "source": [
        "def create_model():\n",
        "  pretrained_model = tf.keras.applications.Xception(input_shape=[*IMAGE_SIZE, 3], include_top=False)\n",
        "  pretrained_model.trainable = True\n",
        "  model = tf.keras.Sequential([\n",
        "    pretrained_model,\n",
        "    tf.keras.layers.GlobalAveragePooling2D(),\n",
        "    tf.keras.layers.Dense(5, activation='softmax')\n",
        "  ])\n",
        "  model.compile(\n",
        "    optimizer='adam',\n",
        "    loss = 'categorical_crossentropy',\n",
        "    metrics=['accuracy']\n",
        "  )\n",
        "  return model\n",
        "\n",
        "with tpu_strategy.scope(): # creating the model in the TPUStrategy scope means we will train the model on the TPU\n",
        "  model = create_model()\n",
        "model.summary()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dMfenMQcxAAb"
      },
      "source": [
        "## Training"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "N0wDLhZOdjAA"
      },
      "source": [
        "Calculate the number of images in each dataset. Rather than actually load the data to do so (expensive), we rely on hints in the filename. This is used to calculate the number of batches per epoch.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "M-ID7vP5mIKs"
      },
      "outputs": [],
      "source": [
        "def count_data_items(filenames):\n",
        "  # The number of data items is written in the name of the .tfrec files, i.e. flowers00-230.tfrec = 230 data items\n",
        "  n = [int(re.compile(r\"-([0-9]*)\\.\").search(filename).group(1)) for filename in filenames]\n",
        "  return np.sum(n)\n",
        "\n",
        "n_train = count_data_items(train_fns)\n",
        "n_valid = count_data_items(validation_fns)\n",
        "train_steps = count_data_items(train_fns) // batch_size\n",
        "print(\"TRAINING IMAGES: \", n_train, \", STEPS PER EPOCH: \", train_steps)\n",
        "print(\"VALIDATION IMAGES: \", n_valid)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eqYT2OcmdnBu"
      },
      "source": [
        "Calculate and show a learning rate schedule. We start with a fairly low rate, as we're using a pre-trained model and don't want to undo all the fine work put into training it."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iXJXG-Ufdbnu"
      },
      "outputs": [],
      "source": [
        "EPOCHS = 12\n",
        "\n",
        "start_lr = 0.00001\n",
        "min_lr = 0.00001\n",
        "max_lr = 0.00005 * tpu_strategy.num_replicas_in_sync\n",
        "rampup_epochs = 5\n",
        "sustain_epochs = 0\n",
        "exp_decay = .8\n",
        "\n",
        "def lrfn(epoch):\n",
        "  if epoch < rampup_epochs:\n",
        "    return (max_lr - start_lr)/rampup_epochs * epoch + start_lr\n",
        "  elif epoch < rampup_epochs + sustain_epochs:\n",
        "    return max_lr\n",
        "  else:\n",
        "    return (max_lr - min_lr) * exp_decay**(epoch-rampup_epochs-sustain_epochs) + min_lr\n",
        "    \n",
        "lr_callback = tf.keras.callbacks.LearningRateScheduler(lambda epoch: lrfn(epoch), verbose=True)\n",
        "\n",
        "rang = np.arange(EPOCHS)\n",
        "y = [lrfn(x) for x in rang]\n",
        "plt.plot(rang, y)\n",
        "print('Learning rate per epoch:')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Zs2tsV8xebhU"
      },
      "source": [
        "Actually train the model. While the first epoch will be quite a bit slower as we must XLA-compile the execution graph and load the data, later epochs should complete in ~5s."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7Qv8rC4aVOFB"
      },
      "outputs": [],
      "source": [
        "history = model.fit(training_dataset, validation_data=validation_dataset,\n",
        "                    steps_per_epoch=train_steps, epochs=EPOCHS, callbacks=[lr_callback])\n",
        "\n",
        "final_accuracy = history.history[\"val_accuracy\"][-5:]\n",
        "print(\"FINAL ACCURACY MEAN-5: \", np.mean(final_accuracy))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VngeUBIdyJ1T"
      },
      "outputs": [],
      "source": [
        "def display_training_curves(training, validation, title, subplot):\n",
        "  ax = plt.subplot(subplot)\n",
        "  ax.plot(training)\n",
        "  ax.plot(validation)\n",
        "  ax.set_title('model '+ title)\n",
        "  ax.set_ylabel(title)\n",
        "  ax.set_xlabel('epoch')\n",
        "  ax.legend(['training', 'validation'])\n",
        "\n",
        "plt.subplots(figsize=(10,10))\n",
        "plt.tight_layout()\n",
        "display_training_curves(history.history['accuracy'], history.history['val_accuracy'], 'accuracy', 211)\n",
        "display_training_curves(history.history['loss'], history.history['val_loss'], 'loss', 212)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "o19yp4TxqmgI"
      },
      "source": [
        "Accuracy goes up and loss goes down. Looks good!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MKFMWzh0Yxsq"
      },
      "source": [
        "## Predictions\n",
        "\n",
        "Let's look at some our model's predictions next to the original images. We'll show 9 images from the validation set."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "S3aW9KYw26f7"
      },
      "outputs": [],
      "source": [
        "def flower_title(label, prediction):\n",
        "  # Both prediction (probabilities) and label (one-hot) are arrays with one item per class.\n",
        "  class_idx = np.argmax(label, axis=-1)\n",
        "  prediction_idx = np.argmax(prediction, axis=-1)\n",
        "  if class_idx == prediction_idx:\n",
        "    return f'{CLASSES[prediction_idx]} [correct]', 'black'\n",
        "  else:\n",
        "    return f'{CLASSES[prediction_idx]} [incorrect, should be {CLASSES[class_idx]}]', 'red'\n",
        "\n",
        "def get_titles(images, labels, model):\n",
        "  predictions = model.predict(images)\n",
        "  titles, colors = [], []\n",
        "  for label, prediction in zip(classes, predictions):\n",
        "    title, color = flower_title(label, prediction)\n",
        "    titles.append(title)\n",
        "    colors.append(color)\n",
        "  return titles, colors\n",
        "\n",
        "validation_viz_iterator = get_dataset_iterator(validation_dataset, 9)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "71fHvf6sjIgX"
      },
      "outputs": [],
      "source": [
        "# Re-run this cell to show a new batch of images\n",
        "images, classes = next(validation_viz_iterator)\n",
        "titles, colors = get_titles(images, classes, model)\n",
        "display_nine_flowers(images, titles, colors)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kIC5KKcJOB_A"
      },
      "source": [
        "## Save and re-loading our trained model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ivTV2qPoOB_B"
      },
      "outputs": [],
      "source": [
        "# We can save our model with:\n",
        "model.save('model.h5')\n",
        "# and reload it with:\n",
        "reloaded_model = tf.keras.models.load_model('model.h5')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZYxyOmXBjA5C"
      },
      "outputs": [],
      "source": [
        "# Re-run this cell to show a new batch of images\n",
        "images, classes = next(validation_viz_iterator)\n",
        "titles, colors = get_titles(images, classes, reloaded_model)\n",
        "display_nine_flowers(images, titles, colors)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a_rjVo-RAoYd"
      },
      "source": [
        "## Next steps\n",
        "\n",
        "More TPU/Keras examples include:\n",
        "- [Shakespeare in 5 minutes with Cloud TPUs and Keras](https://colab.research.google.com/github/tensorflow/tpu/blob/master/tools/colab/shakespeare_with_tpu_and_keras.ipynb)\n",
        "- [Fashion MNIST with Keras and TPUs](https://colab.research.google.com/github/tensorflow/tpu/blob/master/tools/colab/fashion_mnist.ipynb)\n",
        "\n",
        "We'll be sharing more examples of TPU use in Colab over time, so be sure to check back for additional example links, or [follow us on Twitter @GoogleColab](https://twitter.com/googlecolab)."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "provenance": [],
      "private_outputs": true,
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "gpuClass": "standard",
    "accelerator": "TPU"
  },
  "nbformat": 4,
  "nbformat_minor": 0
}