TPU_in_Colab.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Hello, TPU in Colab",
      "version": "0.3.2",
      "provenance": [],
      "collapsed_sections": [],
      "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": [
        "<a href=\"https://colab.research.google.com/github/alxfed/time-to-test-new-tensorflow/blob/master/TPU_in_Colab.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "metadata": {
        "colab_type": "text",
        "id": "_pQCOmISAQBu"
      },
      "cell_type": "markdown",
      "source": [
        "# Testing out the TPU connection\n",
        "\n",
        "First, you'll need to enable TPUs for the notebook.\n",
        "\n",
        "Navigate to Edit→Notebook Settings, and select TPU from the Hardware Accelerator drop-down (you can also access Notebook Settings via the command palette: cmd/ctrl-shift-P).\n",
        "\n",
        "Next, we'll check that we can connect to the TPU."
      ]
    },
    {
      "metadata": {
        "colab_type": "code",
        "id": "71iSWtsXe36x",
        "outputId": "ecb47d9e-3218-45c7-a774-1acd2af49fb0",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 235
        }
      },
      "cell_type": "code",
      "source": [
        "import os\n",
        "import pprint\n",
        "import tensorflow as tf\n",
        "\n",
        "if 'COLAB_TPU_ADDR' not in os.environ:\n",
        "  print('ERROR: Not connected to a TPU runtime; please see the first cell in this notebook for instructions!')\n",
        "else:\n",
        "  tpu_address = 'grpc://' + os.environ['COLAB_TPU_ADDR']\n",
        "  print ('TPU address is', tpu_address)\n",
        "\n",
        "  with tf.Session(tpu_address) as session:\n",
        "    devices = session.list_devices()\n",
        "    \n",
        "  print('TPU devices:')\n",
        "  pprint.pprint(devices)"
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "TPU address is grpc://10.47.178.186:8470\n",
            "TPU devices:\n",
            "[_DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:CPU:0, CPU, -1, 9823645202449288622),\n",
            " _DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 17179869184, 14572425033780392324),\n",
            " _DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:TPU:0, TPU, 17179869184, 2476076947722104267),\n",
            " _DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:TPU:1, TPU, 17179869184, 6477494660992085440),\n",
            " _DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:TPU:2, TPU, 17179869184, 1442762670624733267),\n",
            " _DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:TPU:3, TPU, 17179869184, 8324637985276799343),\n",
            " _DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:TPU:4, TPU, 17179869184, 4976842286325119101),\n",
            " _DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:TPU:5, TPU, 17179869184, 1587310720398405884),\n",
            " _DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:TPU:6, TPU, 17179869184, 16857056776530808963),\n",
            " _DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:TPU:7, TPU, 17179869184, 17346805056559747219),\n",
            " _DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:TPU_SYSTEM:0, TPU_SYSTEM, 17179869184, 7057105281355961679)]\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "metadata": {
        "colab_type": "text",
        "id": "Jkh7cEWRAEA-"
      },
      "cell_type": "markdown",
      "source": [
        "If the cell above reports an error, make sure that you have enabled TPU support in the notebook settings. (Edit menu → Notebook settings)\n",
        "\n",
        "Now, let's try a simple computation."
      ]
    },
    {
      "metadata": {
        "colab_type": "code",
        "id": "0iqSZvc6AX1m",
        "outputId": "ac8f8658-b5ae-4b1b-d5f5-7581deb70288",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 185
        }
      },
      "cell_type": "code",
      "source": [
        "import numpy as np\n",
        "\n",
        "def add_op(x, y):\n",
        "  return x + y\n",
        "  \n",
        "x = tf.placeholder(tf.float32, [10,])\n",
        "y = tf.placeholder(tf.float32, [10,])\n",
        "tpu_ops = tf.contrib.tpu.rewrite(add_op, [x, y])\n",
        "  \n",
        "session = tf.Session(tpu_address)\n",
        "try:\n",
        "  print('Initializing...')\n",
        "  session.run(tf.contrib.tpu.initialize_system())\n",
        "  print('Running ops')\n",
        "  print(session.run(tpu_ops, {x: np.arange(10), y: np.arange(10)}))\n",
        "finally:\n",
        "  # For now, TPU sessions must be shutdown separately from\n",
        "  # closing the session.\n",
        "  session.run(tf.contrib.tpu.shutdown_system())\n",
        "  session.close()"
      ],
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "\n",
            "WARNING: The TensorFlow contrib module will not be included in TensorFlow 2.0.\n",
            "For more information, please see:\n",
            "  * https://github.com/tensorflow/community/blob/master/rfcs/20180907-contrib-sunset.md\n",
            "  * https://github.com/tensorflow/addons\n",
            "If you depend on functionality not listed there, please file an issue.\n",
            "\n",
            "Initializing...\n",
            "Running ops\n",
            "[array([ 0.,  2.,  4.,  6.,  8., 10., 12., 14., 16., 18.], dtype=float32)]\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "metadata": {
        "colab_type": "text",
        "id": "nhXwaCNWBK2n"
      },
      "cell_type": "markdown",
      "source": [
        "# TPU FLOPs\n",
        "\n",
        "Finally, we'll try a small test of floating point computations (floating point operations per seconds. (The units are FLOPS: floating point operations per second.)"
      ]
    },
    {
      "metadata": {
        "colab_type": "code",
        "id": "llcFb_P_BNPM",
        "outputId": "8a00e871-a2e3-463e-e705-1f24dd82b180",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 138
        }
      },
      "cell_type": "code",
      "source": [
        "N = 4096\n",
        "COUNT = 100\n",
        "import time\n",
        "\n",
        "def flops():\n",
        "  x = tf.random_uniform([N, N])\n",
        "  y = tf.random_uniform([N, N])\n",
        "  def _matmul(x, y):\n",
        "    return tf.tensordot(x, y, axes=[[1], [0]]), y\n",
        "\n",
        "  return tf.reduce_sum(\n",
        "    tf.contrib.tpu.repeat(COUNT, _matmul, [x, y])\n",
        "  )\n",
        "  \n",
        "tpu_ops = tf.contrib.tpu.batch_parallel(flops, [], num_shards=8)\n",
        "  \n",
        "session = tf.Session(tpu_address)\n",
        "try:\n",
        "  print('Warming up...')\n",
        "  session.run(tf.contrib.tpu.initialize_system())\n",
        "  session.run(tpu_ops)\n",
        "  print('Profiling')\n",
        "  start = time.time()\n",
        "  session.run(tpu_ops)\n",
        "  end = time.time()\n",
        "  elapsed = end - start\n",
        "  print(elapsed, 'TFlops: {:.2f}'.format(1e-12 * 8 * COUNT * 2*N*N*N / elapsed))\n",
        "finally:\n",
        "  session.run(tf.contrib.tpu.shutdown_system())\n",
        "  session.close()"
      ],
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/control_flow_ops.py:423: colocate_with (from tensorflow.python.framework.ops) is deprecated and will be removed in a future version.\n",
            "Instructions for updating:\n",
            "Colocations handled automatically by placer.\n",
            "Warming up...\n",
            "Profiling\n",
            "0.6735258102416992 TFlops: 163.25\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "metadata": {
        "colab_type": "text",
        "id": "a_rjVo-RAoYd"
      },
      "cell_type": "markdown",
      "source": [
        "# Next steps\n",
        "\n",
        "More involved 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",
        "- [Shakespeare in 5 minutes with Cloud TPUs via TPUEstimator](https://colab.research.google.com/github/tensorflow/tpu/blob/master/tools/colab/shakespeare_with_tpuestimator.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).\n",
        "\n",
        "Meanwhile, you can check out the [TPUEstimator documentation on TensorFlow.org](https://www.tensorflow.org/api_docs/python/tf/contrib/tpu/TPUEstimator). TPUEstimator is an easy way to update models to take advantage of TPU acceleration."
      ]
    }
  ]
}