GPUvsTPU.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "GPUvsTPU.ipynb",
      "version": "0.3.2",
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "TPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "[View in Colaboratory](https://colab.research.google.com/github/zaidalyafeai/Notebooks/blob/master/GPUvsTPU.ipynb)"
      ]
    },
    {
      "metadata": {
        "id": "QueVjeESsyKe",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "# Why TPUs ? "
      ]
    },
    {
      "metadata": {
        "id": "5moeHHv4shGw",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "TPUs are tensor processing units developed by Google to  accelerate operations on a Tensorflow Graph. Each TPU packs up to 180 teraflops of floating-point performance and 64 GB of high-bandwidth memory onto a single board. Here is a comparions between TPUs and Nvidia GPUs. The y axis represents # images per seconds and the x axis is different models.\n",
        "\n",
        "<img src=\"https://cdn-images-1.medium.com/max/800/1*tVHGjJHJrhKaKECT3Z4CIw.png\" alt=\"Drawing\" style=\"width: 150px;\"/>"
      ]
    },
    {
      "metadata": {
        "id": "_SXoMcRs8aRs",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "# Experiement \n",
        "\n",
        "TPUs were only available on Google cloud but now they are available for free in Colab. We will be comparing TPU vs GPU here on colab using mnist dataset. We will compare the time of each step and epoch against different batch sizes. "
      ]
    },
    {
      "metadata": {
        "id": "4ECTupP8warH",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "# Downoad MNIST "
      ]
    },
    {
      "metadata": {
        "id": "oX7DOjhUlCLb",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "import tensorflow as tf\n",
        "import os\n",
        "import numpy as np\n",
        "from tensorflow.keras.utils import to_categorical\n",
        "\n",
        "def get_data():\n",
        "\n",
        "  #Load mnist data set\n",
        "  (x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()\n",
        "\n",
        "  x_train = x_train.astype('float32') / 255\n",
        "  x_test = x_test.astype('float32') / 255\n",
        "\n",
        "  x_train = np.expand_dims(x_train, 3)\n",
        "  x_test = np.expand_dims(x_test, 3)\n",
        "\n",
        "  y_train = to_categorical(y_train)\n",
        "  y_test  = to_categorical(y_test)\n",
        "\n",
        "  return x_train, y_train, x_test, y_test "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "AtEJD_s1wdty",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "# Basic CNN"
      ]
    },
    {
      "metadata": {
        "id": "IaZZ2OwmwhKQ",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "Note that since we need to run the code on TPU we need to do more work. We need to specify the address of the TPU and tell tensorflow to run the model on the TPU cluster "
      ]
    },
    {
      "metadata": {
        "id": "cUYn3VomnQDL",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "from tensorflow.contrib.tpu.python.tpu import keras_support\n",
        "\n",
        "def get_model(tpu = False):\n",
        "  model = tf.keras.Sequential()\n",
        "\n",
        "  #add layers to the model \n",
        "  model.add(tf.keras.layers.Conv2D(filters=64, kernel_size=2, padding='same', activation='relu', input_shape=(28,28,1))) \n",
        "  model.add(tf.keras.layers.MaxPooling2D(pool_size=2))\n",
        "  model.add(tf.keras.layers.Dropout(0.3))\n",
        "\n",
        "  model.add(tf.keras.layers.Conv2D(filters=32, kernel_size=2, padding='same', activation='relu'))\n",
        "  model.add(tf.keras.layers.MaxPooling2D(pool_size=2))\n",
        "  model.add(tf.keras.layers.Dropout(0.3))\n",
        "\n",
        "  model.add(tf.keras.layers.Flatten())\n",
        "  model.add(tf.keras.layers.Dense(256, activation='relu'))\n",
        "  model.add(tf.keras.layers.Dropout(0.5))\n",
        "  model.add(tf.keras.layers.Dense(10, activation='softmax'))\n",
        "\n",
        "  #compile the model \n",
        "  model.compile(loss='categorical_crossentropy',\n",
        "               optimizer='adam',\n",
        "               metrics=['accuracy'])\n",
        "\n",
        "  #flag to run on tpu \n",
        "  if tpu:\n",
        "    tpu_grpc_url = \"grpc://\"+os.environ[\"COLAB_TPU_ADDR\"]\n",
        "    \n",
        "    #connect the TPU cluster using the address \n",
        "    tpu_cluster_resolver = tf.contrib.cluster_resolver.TPUClusterResolver(tpu_grpc_url)\n",
        "    \n",
        "    #run the model on different clusters \n",
        "    strategy = keras_support.TPUDistributionStrategy(tpu_cluster_resolver)\n",
        "    \n",
        "    #convert the model to run on tpu \n",
        "    model = tf.contrib.tpu.keras_to_tpu_model(model, strategy=strategy)\n",
        "  return model"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "cSoBDg4PwylQ",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "#GPU vs TPU\n"
      ]
    },
    {
      "metadata": {
        "id": "yluf1xqjsILa",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "x_train, y_train, x_test, y_test = get_data()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "Wn9zXSBUw8m1",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "Each time you want to run the model on TPU make sure to set the tpu flag and change the enviornment runtime via  Edit> Notebook Setting > Hardware Accelerator > TPU and then click save. "
      ]
    },
    {
      "metadata": {
        "id": "4vAM7pBPxVbm",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "#set tpu = True if you want to run the model on TPU\n",
        "model = get_model(tpu = False)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "_m67tWDhnm7f",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 139
        },
        "outputId": "59645203-25f4-4e4f-87a2-72da5fa1cf48"
      },
      "cell_type": "code",
      "source": [
        "model.fit(x_train,\n",
        "         y_train,\n",
        "         batch_size=1024,\n",
        "         epochs=10,\n",
        "         validation_data=(x_test, y_test))"
      ],
      "execution_count": 18,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Train on 60000 samples, validate on 10000 samples\n",
            "Epoch 1/3\n",
            "60000/60000 [==============================] - 2s 38us/step - loss: 0.1639 - acc: 0.9513 - val_loss: 0.0677 - val_acc: 0.9752\n",
            "Epoch 2/3\n",
            "60000/60000 [==============================] - 2s 35us/step - loss: 0.1345 - acc: 0.9573 - val_loss: 0.0552 - val_acc: 0.9808\n",
            "Epoch 3/3\n",
            "60000/60000 [==============================] - 2s 40us/step - loss: 0.1189 - acc: 0.9619 - val_loss: 0.0443 - val_acc: 0.9848\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "metadata": {
        "id": "QGof6K46zXfq",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "# Benchmarks \n",
        "\n",
        "Note that TPU setup takes some time when compiling the model and distributing the data in the clusters, so the first epoch will take alonger time. I only reported the time for the later epochs. I calculated the average time accross different epochs."
      ]
    },
    {
      "metadata": {
        "id": "4cbKs72g00sQ",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### Epoch Time ($s$)"
      ]
    },
    {
      "metadata": {
        "id": "QNh64VMDz1Ks",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "$$\\left[\\begin{array}{c|c|c}  \n",
        " \\textbf{Batch Size} & \\textbf{GPU} & \\textbf{TPU} \\\\\n",
        " 256 & 6s & 6s\\\\  \n",
        " 512 & 5s & 3s\\\\\n",
        " 1024 & 4s & 2s\\\\\n",
        "\\end{array}\\right]$$"
      ]
    },
    {
      "metadata": {
        "id": "Q8eMm1GD1Mu5",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### Step Time ($\\mu s$)"
      ]
    },
    {
      "metadata": {
        "id": "q1hElmjr05Ah",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "$$\\left[\\begin{array}{c|c|c}  \n",
        " \\textbf{Batch Size} & \\textbf{GPU} & \\textbf{TPU} \\\\\n",
        " 256 & 94 \\mu s & 97 \\mu s\\\\  \n",
        " 512 & 82 \\mu  s& 58 \\mu s \\\\\n",
        " 1024 & 79 \\mu s & 37 \\mu s\\\\\n",
        "\\end{array}\\right]$$"
      ]
    },
    {
      "metadata": {
        "id": "J6eOKfyW38rN",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "# References\n",
        "\n",
        "\n",
        "\n",
        "*   https://qiita.com/koshian2/items/25a6341c035e8a260a01\n",
        "*   https://medium.com/tensorflow/hello-deep-learning-fashion-mnist-with-keras-50fcff8cd74a\n",
        "*   https://blog.riseml.com/benchmarking-googles-new-tpuv2-121c03b71384\n",
        "*   https://cloudplatform.googleblog.com/2018/02/Cloud-TPU-machine-learning-accelerators-now-available-in-beta.html\n",
        "\n"
      ]
    }
  ]
}