CPU_vs_GPU_vs_TPU_Inference_Timing.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyORjSUqNkfyQ/hzibzWDlNm",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "gpuClass": "standard",
    "accelerator": "TPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/Ekhao/CPUvsGPUvsTPUInferenceTiming/blob/main/CPU_vs_GPU_vs_TPU_Inference_Timing.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# CPU vs GPU vs TPU Inference Time\n",
        "This colab notebook aims to provide a comparison between the inference performance of CPUs, GPUs and TPUs.\n",
        "\n",
        "Unfortunately it is not possible to dynamically switch between computing instances on Google Colab. Therefore this notebook must be run once for each type of computing platform. Some instructions may be specific to one platform e.g. the TPU. These instructions are marked as such and put in standalone code blocks."
      ],
      "metadata": {
        "id": "erqS7OSEfNzl"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Imports"
      ],
      "metadata": {
        "id": "Arfq8m-hfg_9"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6abVElHzeOx7"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "import tensorflow_datasets as tfds"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## CPU Specifics\n",
        "We don't need to run any specific code to run inference using the CPU. However, we can check what cpu we are using by using the \"lscpu\" command."
      ],
      "metadata": {
        "id": "TzZhluX7ldX0"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!lscpu\n",
        "strategy = tf.distribute.get_strategy()"
      ],
      "metadata": {
        "id": "dg_Drg2nlvM7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## GPU Specifics\n",
        "We also do not need to run any special code to run inference on a GPU. We can again however, check the GPU that we are using by using the \"nvidia-smi\" command."
      ],
      "metadata": {
        "id": "HO3QOVXcmlMI"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!nvidia-smi\n",
        "strategy = tf.distribute.get_strategy()"
      ],
      "metadata": {
        "id": "kg33KZMwm-r7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## TPU Specifics\n",
        "Unlike the CPU and GPU, the TPU requires some code to set up. This is likely because it is a newer platform which does not have as streamlined of an interface as the CPU and GPU. \n",
        "\n",
        "We also print a list of TPU devices that we are using."
      ],
      "metadata": {
        "id": "P6zMmjuTpq6e"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "try:\n",
        "  tpu = tf.distribute.cluster_resolver.TPUClusterResolver(tpu=\"\")\n",
        "  print(\"TPU Available\")\n",
        "except ValueError:\n",
        "  tpu = None\n",
        "if tpu:\n",
        "  tf.config.experimental_connect_to_cluster(tpu)\n",
        "  tf.tpu.experimental.initialize_tpu_system(tpu)\n",
        "  strategy = tf.distribute.TPUStrategy(tpu)"
      ],
      "metadata": {
        "id": "YWCLzyeJqMg-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Set up the a model\n",
        "For this experiment we choose to use the ResNetV2 Neural Network. This model can be loaded easily through keras."
      ],
      "metadata": {
        "id": "VNLv6gAFggbl"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "with strategy.scope():\n",
        "  model = tf.keras.applications.ResNet152V2(\n",
        "      include_top=True,\n",
        "      weights=\"imagenet\",\n",
        "  )"
      ],
      "metadata": {
        "id": "vWpcWSlWg1f7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Set up a Dataset\n",
        "ResNet152V2 is a model meant to be processing image data. It is trained on Imagenet - but as we don't really care about accurate predictions we can use any available dataset. \n",
        "\n",
        "We choose to use the \"tf_flowers\" as it is available in google cloud storage for free. Note that storage in google cloud storage is required for TPUs."
      ],
      "metadata": {
        "id": "fyvLABmnie3S"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "data = tfds.load(\"tf_flowers\", split=\"train\", as_supervised=True, try_gcs=True)"
      ],
      "metadata": {
        "id": "8ky3jyT3iyy-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "We need to preprocess the dataset a bit for two reasons:\n",
        "\n",
        "\n",
        "*   The images of tf_flowers are of different sizes. We need to resize them all to 224x224 pixels to be used with the ResNet152V2 model.\n",
        "*   A tensorflow dataset can be batched, cached and prefetched to optimize memory latency. To provide a fair comparison between CPUs, GPUs and TPUs we apply these optimizations so that each processing unit can work as fast as possible.\n",
        "\n"
      ],
      "metadata": {
        "id": "FkEUFdeci88S"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# A thin wrapper function that takes an entry of a tensorflow dataset loaded \"as_supervised\" and resizes the image in the entry\n",
        "def resize_image(image, label):\n",
        "  return tf.image.resize(image, [224,224]), label\n",
        "\n",
        "\n",
        "data = data.map(resize_image, num_parallel_calls=tf.data.AUTOTUNE).repeat(3).batch(128).cache().prefetch(tf.data.AUTOTUNE)"
      ],
      "metadata": {
        "id": "R5YbYqmpwDY7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Run Inference\n",
        "In this block we make the ResNet152V2 model run inference on the tf_flowers dataset.\n",
        "\n",
        "For fairness it makes sense to run this twice to reduce the time spent on overhead for first loading."
      ],
      "metadata": {
        "id": "HWo9FZfpjy17"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "model.predict(data)"
      ],
      "metadata": {
        "id": "zg3rMcETYk_A"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}