TPU_MachineLearning (2).ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "TPU_MachineLearning.ipynb",
      "version": "0.3.2",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "metadata": {
        "id": "VCGujgRo19ld",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### Lecture Structure\n",
        "- Q&A\n",
        "- Why make a TPU? (10 minutes)\n",
        "- Short demo of a matrix operation (5 minutes)\n",
        "- Q&A \n",
        "- TPU Explained (technical) (10 minutes)\n",
        "- TPU Use cases (5 minutes) \n",
        "- Benchmark Demo #1 (Simple Matrix Add Operation) for CPU vs TPU (10 minutes)\n",
        "- Q&A \n",
        "- Benchmark Demo #2 (Neural Network for Image Classificatio) for GPU vs TPU (10 minutes)\n",
        "- Demo #3 Training a text generation on a TPU (10 minutes)\n",
        "- Rap, probably. Lol\n",
        "\n"
      ]
    },
    {
      "metadata": {
        "id": "W-egUimiTrzt",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## What is a Tensor Processing Unit?\n",
        "\n",
        "![alt text](http://www.cdrinfo.com/images/uploaded/Google_TPU_3.jpg)\n",
        "\n",
        "- A Tensor Processing Unit (TPU) is a custom computer processing chip designed by Google\n",
        "\n",
        "![alt text](https://i.imgur.com/dMx3Ilw.png)\n",
        "\n",
        "- They've been using TPUs in their data centers since 2015\n",
        "\n",
        "![alt text](https://i.imgur.com/TnpIdxd.png)\n",
        "\n",
        "- They designed them specifically for machine learning applications\n",
        "- They use them for Google Translate, Photos, Search Assistant, Gmail, Cloud, etc. \n",
        "\n",
        "## Why did they make their own chip? \n",
        "\n",
        "![alt text](https://i.imgur.com/N1eOy9m.png)\n",
        "\n",
        "- There has been a lot of progress in machine learning in the past few years\n",
        "\n",
        "![alt text](https://image.slidesharecdn.com/20170222mldlintroduction-170222094304/95/machine-learning-deep-learning-and-data-analysis-introduction-38-638.jpg?cb=1496630353)\n",
        "\n",
        "- Neural Networks in particular almost always outperform other machine learning models if given enough data & compute\n",
        "- Neural networks require a lot of compute! It's called deep learning.\n",
        "\n",
        "![alt text](https://image.slidesharecdn.com/2015scalaworld-150927123309-lva1-app6892/95/neural-network-as-a-function-12-638.jpg?cb=1443357296)\n",
        "\n",
        "![alt text](https://leonardoaraujosantos.gitbooks.io/artificial-inteligence/content/Images/Linear_1.jpg)\n",
        "\n",
        "- Neural networks are just a series of matrix operations applied to input data\n",
        "- And if theres a lot of data to input, thats a lot of matrix operations to compute\n",
        "- Like a lot. Giant, giant matrices full of numbers all being multiplied in parallel\n",
        "-Most of the math is just 'multiply a bunch of numbers, and add the results' \n",
        "- We can connect these two together in a single operation called multiply-accumulate (MAC). \n",
        "- And if we don’t need to do anything else, we can multiply-accumulate really, really fast.\n",
        "\n",
        "![alt text](https://c1.staticflickr.com/2/1640/25046013104_68059057ab_b.jpg)\n",
        "\n",
        "![alt text](https://cdn-images-1.medium.com/max/800/1*OZWu6AW9nI9HXb4g1XFCQA.png)\n",
        "\n",
        "- A CPU is a scalar machine, which means it processes instructions one step at a time. \n",
        "- CPUs can peform these matrix operations pretty well.\n",
        "- But Moore's Law is coming to an end. \n",
        "\n",
        "![alt text](https://www.carestream.com/blog/wp-content/uploads/2015/09/CSH_CPU-GPU_Illustration.png)\n",
        "\n",
        "- A CPU is composed of just a few cores with lots of cache memory that can handle a few software threads at a time\n",
        "- Luckily, GPUs (Graphics Processing Unit) can perform matrix operations orders of magnitude better than CPUs.\n",
        "\n",
        "![alt text](https://www.researchgate.net/profile/Abel_Paz/publication/231167191/figure/fig2/AS:300437820461057@1448641366026/Comparison-of-CPU-versus-GPU-architecture.png)\n",
        "\n",
        " - A GPU is composed of hundreds of cores that can handle thousands of threads simultaneously.\n",
        "- Thats because GPUs were designed for 3d game rendering, which often involves parallel operations\n",
        "-The ability of a GPU with 100+ cores to process thousands of threads can accelerate some software by 100x over a CPU alone. \n",
        "- What’s more, the GPU achieves this acceleration while being more power- and cost-efficient than a CPU.\n",
        "- So when neural networks run on GPUs, they run much faster than on CPUs\n",
        "\n",
        "![alt text](http://share.opsy.st/54176dc2ec16c-Vivante-Sept-Fig1.jpg)\n",
        "![alt text](https://cdn-images-1.medium.com/max/600/1*Yx9XF4H4spE8Bm8XVY1bYA.png)\n",
        "\n",
        "- A GPU is a vector machine. You can give it a long list of data — a 1D vector — and run computations on the entire list at the same time. \n",
        "- This way, we can perform more computations per second, but we have to perform the same computation on a vector of data in parallel. \n",
        "- GPUs are general purpose chips. They don't just perform matrix operations, they can really do any kind of computation.\n",
        "- GPUs are optimized for taking huge batches of data and performing the same operation over and over very quickly\n"
      ]
    },
    {
      "metadata": {
        "id": "ujHNgy3sYzBj",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## The Power of Linear Algebra\n",
        "\n",
        "- Lets perform a simple matrix operation on the GPU using the MXNet machine learning\n",
        "- This is basically CUDA (GPU Programming) MXNet's NDArray is a thin wrapper on top of CUDA constructs. Awesome. \n",
        "- Tensorflow uses lots of CUDA as well. \n",
        "\n",
        "![alt text](https://cdn.kastatic.org/googleusercontent/xj8YqV88KB29MEKR5Iq68oUo1h2kFAIAewMsHeWS9-7l0KaB6BI3sOmpfGSCzsVU8z5Evq6QIrwbEAqBnZ5W06g0CQ)"
      ]
    },
    {
      "metadata": {
        "id": "naqcUZaHW48V",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "#1 - MXNet Dependencies"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "0R0MRer6XX_L",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "#2 Basically CUDA time. CUDA is awesomeeeeeeee <3 u Nvidia "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "i4YyHMGAGN2P",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "# Q&A "
      ]
    },
    {
      "metadata": {
        "id": "JpeL3jWPc7kD",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## The TPU\n",
        "\n",
        "![alt text](https://pbs.twimg.com/media/Dcww6cOUQAAgKYb.jpg)\n",
        "\n",
        "- 3 Generations \n",
        "- Google took 15 months for the TPUv1, and that was astonishingly fast for an ASIC.\n",
        "- ASICs are initially expensive, requiring specialized engineers and manufacturing costs that start at around a million dollars. \n",
        "- And they are inflexible: there’s no way to change the chip once it’s finished.\n",
        "- But if you know you’ll be doing one particular job in enough volume, the recurring benefits can make up for the initial drawbacks. \n",
        "- ASICs are generally the fastest and most energy-efficient way to accomplish a task. \n",
        "\n",
        "3 Generation of TPUS. Describe a TPU. Then Simple CPU VS TPU Add benchmark. Then Japanese MLP GPU vs TPU benchmark. Then Tensorflow Shakespeare. \n",
        "\n",
        "![alt text](https://cdn-images-1.medium.com/max/800/1*dW4_Z8wsMcoS44dtKNacZQ.gif)\n",
        "\n",
        "-The data of a neural network is arranged in a matrix, a 2D vector.\n",
        "- So, Google decided they needed to build a a matrix machine (The tensor processing unit or TPU)\n",
        "- And they really only care about multiply-accumulate, so they prioritized that over other instructions that a processor would normally support.\n",
        "- We’ll devote most of our chip to the MACs that perform matrix multiplication, and mostly ignore other operations.²\n",
        "- Google wanted to design a chip specifically for the matrix operations that neural networks require so that it would run them even more efficiently.\n",
        "\n",
        "![alt text](https://cloud.google.com/tpu/docs/images/tpu--sys-arch3.png)\n",
        "\n",
        "- TPU hardware is comprised of four independent chips.\n",
        "- The following block diagram describes the components of a single chip. \n",
        "- Each chip consists of two compute cores called Tensor Cores. \n",
        "- A Tensor Core consists of scalar, vector and matrix units (MXU). \n",
        "- In addition, 8 GB of on-chip memory (HBM) is associated with each Tensor Core.\n",
        "- The bulk of the compute horsepower in a Cloud TPU is provided by the MXU. \n",
        "- Each MXU is capable of performing 16K multiply-accumulate operations in each cycle. \n",
        "- While the MXU's inputs and outputs are 32-bit floating point values, the MXU performs multiplies at reduced bfloat16 precision. \n",
        "- Bfloat16 is a 16-bit floating point representation that provides better training and model accuracy than the IEEE half-precision representation.\n",
        "-From a software perspective, each of the 8 cores on a Cloud TPU can execute user computations (XLA ops) independently. \n",
        "- High-bandwidth interconnects allow the chips to communicate directly with each other.\n",
        "\n",
        "![alt text](https://cdn-images-1.medium.com/max/600/1*d7Lg4cYdQO2kxt9nSmyj-g.png)\n",
        "\n",
        "- XLA is an experimental JIT compiler for the backend of Tensorflow. \n",
        "- It turns your TF graph into linear algebra, and it has backends of its own to run on CPUs, GPUs, or TPUs\n",
        "- proprietary \n",
        "\n",
        "## The Systolic Array\n",
        "\n",
        "- The way to achieve that matrix performance is through a piece of architecture called a systolic array.\n",
        "- This is the interesting bit, and it’s why a TPU is performant. \n",
        "- A systolic array is a kind of hardware algorithm, and it describes a pattern of cells on a chip that computes matrix multiplication. \n",
        "- “Systolic” describes how data moves in waves across the chip, like the beating of a human heart.\n",
        "\n",
        "![alt text](https://cdn-images-1.medium.com/max/800/1*umH-qhj3j1Z1k35uBkXM8g.png)\n",
        "\n",
        "- On a TPU, For 2x2 inputs, each term in the output is the sum of two products. \n",
        "- No product is reused, but the individual terms are.\n",
        "- We’ll implement this by building a 2x2 grid. (It’s actually a grid, not just an abstraction — hardware is fun like that).\n",
        "- Note that 2x2 is a toy example, and the full size MXU is a monstrous 128x128.\n",
        "- Let’s say AB/CD represents our activations and EF/GH our weights. For our array, we first load up the weights like so\n",
        "\n",
        "![alt text](https://cdn-images-1.medium.com/max/800/1*P83xXpFMjgLkAXf9i4uC9g.png)\n",
        "\n",
        "- Our activations go into an input queue, which for our example will go on the left of each row.\n",
        "\n",
        "![alt text](https://cdn-images-1.medium.com/max/800/1*PyZRdGouPq26ON8G1vo7gQ.png)\n",
        "\n",
        "- Every cycle of the clock, each cell will execute the following steps, all in parallel:\n",
        "\n",
        "1. Multiply our weight and the activation coming in from the left. If there’s no cell on the left, take from the input queue.\n",
        "2. Add that product to the partial sum coming in from above. If there’s no cell above, the partial sum from above is zero.\n",
        "3. Pass the activation to the cell to the right. If there’s no cell on the right, throw away the activation.\n",
        "4. Pass the partial sum to the cell on the bottom. If there’s no cell on the bottom, collect the partial sum as an output.\n",
        "\n",
        "[Python Example](https://github.com/antonpaquin/SystolicArrayDemo/blob/master/systolic.py)\n",
        "\n",
        "By these rules, you can see the activations will start on the left and move one cell to the right per cycle, and the partial sums will start on top amd move one cell down per cycle.\n",
        "\n",
        "![alt text](https://cdn-images-1.medium.com/max/800/1*b7t8CxJJ8K61BC0NHXjfyg.png )\n",
        "\n",
        "- Thats our data flow. \n",
        "\n",
        "- So one Cycle of many happening in paralellel, for example, looks like this  \n",
        "\n",
        "1. Top left reads A from input queue, multiplies with weight E to produce product AE.\n",
        "2. AE added to partial sum 0 from above, produces partial sum AE.\n",
        "3. Activation A passed to the cell in the top right.\n",
        "4. Partial sum AE passed to the cell in the bottom left.\n",
        "\n",
        "![alt text](https://cdn-images-1.medium.com/max/600/1*_v7REW2VWoVsob9HS7kMFw.gif)\n",
        "\n",
        "- It takes 3n-2 cycles to fully compute the result matrix, whereas a standard sequential solution is n³. \n",
        "\n",
        "![alt text](https://cdn-images-1.medium.com/max/600/1*IyFcYIFIaMwAn90RIJYfMw.png)\n",
        "\n",
        "- We can do this because we’re running 128x128 MAC operations in parallel. \n",
        "- Multipliers are usually big and expensive to implement in hardware, but the high density of systolic arrays lets Google pack 16,384 of them into the MXU. \n",
        "- This ranslates directly into speed training and running your network.\n",
        "- Weights are loaded in much the same way as activations — through the input queue. \n",
        "- We just send a special control signal (red in the above diagram) to tell the array to store weights as they pass by, instead of running MAC operations.\n",
        "- Weights remain in the same processing elements, so we can send an entire batch through before loading a new set, reducing overhead.\n",
        "- That’s it! The rest of the chip is important and worth going over, but the core advantage of the TPU is its MXU — a systolic array matrix multiplication unit.\n",
        "\n",
        "\n",
        "## What are the Use Cases for TPUs?\n",
        "\n",
        "![alt text](https://storage.googleapis.com/gweb-cloudblog-publish/original_images/tpu-6tlel.PNG)\n",
        "\n",
        "CPUs:\n",
        "\n",
        "- Quick prototyping that requires maximum flexibility\n",
        "- Simple models that do not take long to train\n",
        "- Small models with small effective batch sizes\n",
        "- Models that are dominated by custom TensorFlow operations written in C++\n",
        "- Models that are limited by available I/O or the networking bandwidth of the host system\n",
        "\n",
        "\n",
        "GPUs:\n",
        "\n",
        "- Models that are not written in TensorFlow or cannot be written in TensorFlow\n",
        "- Models for which source does not exist or is too onerous to change\n",
        "- Models with a significant number of custom TensorFlow operations that must run at least partially on CPUs\n",
        "- Models with TensorFlow ops that are not available on Cloud TPU (see the list of available TensorFlow ops)\n",
        "- Medium-to-large models with larger effective batch sizes\n",
        "\n",
        "TPUs:\n",
        "\n",
        "- Models dominated by matrix computations\n",
        "- Models with no custom TensorFlow operations inside the main training loop\n",
        "- Models that train for weeks or months\n",
        "- Larger and very large models with very large effective batch sizes\n"
      ]
    },
    {
      "metadata": {
        "id": "DG5ZtKTQgRhL",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "# Benchmark #1 - Simple matrix add operation on CPU vs TPU"
      ]
    },
    {
      "metadata": {
        "id": "9sEUvuISDtO3",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# 3 Dependencies for benchmark 1"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "rK2y3WIAEciE",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "#4 CPU Benchmark"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "lxViRKNJW3K9",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# 5 TPU Benchmark"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "cJ3URGN-W60G",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# 6 TPU Continued"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "QXn751_GGbkG",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "# Q&A "
      ]
    },
    {
      "metadata": {
        "id": "V6LVUTDpuBl0",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## Key TPU Functions in Tensorflow\n",
        "\n",
        "#### The TPU Estimator\n",
        "\n",
        "- Estimators are TensorFlow's model-level abstraction.\n",
        "- Standard Estimators can drive models on CPU and GPUs. You must use tf.contrib.tpu.TPUEstimator to drive a model on TPUs.\n",
        "\n",
        "```\n",
        "my_tpu_estimator = tf.contrib.tpu.TPUEstimator(\n",
        "    model_fn=my_model_fn,\n",
        "    config=tf.contrib.tpu.RunConfig()\n",
        "    use_tpu=False)\n",
        "```\n",
        "#### The TPU Run Configuration\n",
        "\n",
        "\n",
        "\n",
        "```\n",
        "my_tpu_run_config = tf.contrib.tpu.RunConfig(\n",
        "    master=master,\n",
        "    evaluation_master=master,\n",
        "    model_dir=FLAGS.model_dir,\n",
        "    session_config=tf.ConfigProto(\n",
        "        allow_soft_placement=True, log_device_placement=True),\n",
        "    tpu_config=tf.contrib.tpu.TPUConfig(FLAGS.iterations,\n",
        "                                        FLAGS.num_shards),\n",
        ")\n",
        "\n",
        "```\n",
        "#### The Cross Shard Optimizer\n",
        "\n",
        "- When training on a cloud TPU you must wrap the optimizer in a tf.contrib.tpu.CrossShardOptimizer, which uses an allreduce to aggregate gradients and broadcast the result to each shard (each TPU core).\n",
        "\n",
        "```\n",
        "optimizer = tf.train.GradientDescentOptimizer(learning_rate=learning_rate)\n",
        "if FLAGS.use_tpu:\n",
        "  optimizer = tf.contrib.tpu.CrossShardOptimizer(optimizer)\n",
        "```\n",
        "\n",
        "#### Tensor dimensions need to be statically defined at compile time thanks to XLA\n",
        "\n",
        "- During regular Tensorflow execution any unknown shape dimensions are determined dynamically\n",
        "- but XLA requires that all tensor dimensions be statically defined at compile time. All shapes must evaluate to a constant, and not depend on external data, or stateful operations like variables or a random number generator.\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "metadata": {
        "id": "TsfIE1SowK21",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "# Benchmark #2 - Multilayer perceptron for image classification using the CIFAR Dataset (GPU vs TPU)"
      ]
    },
    {
      "metadata": {
        "id": "nNEHDKPOzzXF",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "import tensorflow as tf\n",
        "import tensorflow.keras.backend as K\n",
        "import numpy as np\n",
        "from tensorflow.contrib.tpu.python.tpu import keras_support\n",
        "from tensorflow.keras.layers import Input, Conv2D, BatchNormalization, Activation, AveragePooling2D, Dense, Dropout, Flatten\n",
        "from tensorflow.keras.datasets import cifar10\n",
        "from tensorflow.keras.models import Model, Sequential\n",
        "from tensorflow.keras.utils import to_categorical\n",
        "import os\n",
        "\n",
        "def basic_mlp_module(input, units):\n",
        "    x = Dense(units)(input)\n",
        "    x = BatchNormalization()(x)\n",
        "    x = Activation(\"relu\")(x)\n",
        "    x = Dropout(0.5)(x)\n",
        "    return x\n",
        "\n",
        "def create_mlp_model():\n",
        "    input = Input((32*32*3,))\n",
        "    x = basic_mlp_module(input, 2048)\n",
        "    x = basic_mlp_module(x, 1024)\n",
        "    x = basic_mlp_module(x, 512)\n",
        "    x = basic_mlp_module(x, 256)\n",
        "    x = basic_mlp_module(x, 128)\n",
        "    x = basic_mlp_module(x, 64)\n",
        "    x = basic_mlp_module(x, 32)\n",
        "    x = basic_mlp_module(x, 16)\n",
        "    x = Dense(10, activation=\"softmax\")(x)\n",
        "    return Model(input, x)\n",
        "\n",
        "def main():\n",
        "    # 7 lets make this into TPU Code!!!\n",
        "    K.clear_session()\n",
        "\n",
        "    # CIFAR\n",
        "    (X_train, y_train), (_, _) = cifar10.load_data()\n",
        "    X_train = (X_train / 255.0).reshape(50000, -1)\n",
        "    y_train = to_categorical(y_train)\n",
        "\n",
        "    # Model building\n",
        "    model = create_mlp_model()\n",
        "    model.compile(tf.train.AdamOptimizer(learning_rate=1e-3), loss=\"categorical_crossentropy\", metrics=[\"acc\"])\n",
        "\n",
        "    model.fit(X_train, y_train, batch_size=1024, epochs=10)\n",
        "\n",
        "if __name__ == \"__main__\":\n",
        "    main()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "dVI91_yA1gN9",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### Lets look at their default 'train shakespeare in 5 minutes example'\n",
        "\n",
        "https://colab.research.google.com/github/tensorflow/tpu/blob/master/tools/colab/shakespeare_with_tpu_and_keras.ipynb\n",
        "\n",
        " "
      ]
    }
  ]
}