Getting_Started_with_PyTorch_on_Cloud_TPUs.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Copy of Getting Started with PyTorch on Cloud TPUs",
      "provenance": [],
      "collapsed_sections": [],
      "machine_shape": "hm",
      "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/jonbaer/googlecolab/blob/master/Getting_Started_with_PyTorch_on_Cloud_TPUs.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RKLajLqUni6H"
      },
      "source": [
        "## Getting Started with PyTorch on Cloud TPUs\n",
        "\n",
        "This notebook will show you how to:\n",
        "\n",
        "* Install PyTorch/XLA on Colab, which lets you use PyTorch with TPUs.\n",
        "* Run basic PyTorch functions on TPUs, like creating and adding tensors.\n",
        "* Run PyTorch modules and autograd on TPUs.\n",
        "* Run PyTorch networks on TPUs.\n",
        "\n",
        "PyTorch/XLA is a package that lets PyTorch connect to Cloud TPUs and use TPU cores as devices. Colab provides a free Cloud TPU system (a remote CPU host + four TPU chips with two cores each) and installing PyTorch/XLA only takes a couple minutes. \n",
        "\n",
        "Even though Colab offers eight TPU cores, this notebook only uses one for simplicity. More information about running PyTorch on TPUs can be found on [PyTorch.org](http://pytorch.org/xla/), including how to run PyTorch networks on multiple TPU cores simultaneously. Other Colab notebooks also show how to use multiple TPU cores, including [this one](https://colab.research.google.com/github/pytorch/xla/blob/master/contrib/colab/mnist-training.ipynb#scrollTo=Afwo4H7kSd8P) which trains a network on the MNIST dataset and [this one](https://colab.research.google.com/github/pytorch/xla/blob/master/contrib/colab/resnet18-training.ipynb#scrollTo=_2nL4HmloEyl) which trains a ResNet18 architecture on CIFAR10. \n",
        "\n",
        "These and other Colab notebooks, as well as Google Cloud Platform (GCP) tutorials, can be found [here](https://github.com/pytorch/xla/tree/master/contrib/colab). Check out our [NeurIPS 2019 Fast Neural Style Transfer demo](https://colab.research.google.com/github/pytorch/xla/blob/master/contrib/colab/style_transfer_inference.ipynb#scrollTo=EozMXwIV9iOJ), where you can apply different styles (filters) to your own images!\n",
        "\n",
        "To use PyTorch on Cloud TPUs in your own Colab notebook you can copy this one, or copy the setup cell below and configure your Colab environment to use TPUs. \n",
        "\n",
        "Finally, this notebook is intended for people already familiar with PyTorch, a popular open-source deep learning framework. If you haven't used PyTorch before you might want to review the tutorials at https://pytorch.org/ before continuing.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b3rCVMRazoeB"
      },
      "source": [
        "<h3>  &nbsp;&nbsp;Use Colab Cloud TPU&nbsp;&nbsp; <a href=\"https://cloud.google.com/tpu/\"><img valign=\"middle\" src=\"https://raw.githubusercontent.com/GoogleCloudPlatform/tensorflow-without-a-phd/master/tensorflow-rl-pong/images/tpu-hexagon.png\" width=\"50\"></a></h3>\n",
        "\n",
        "* On the main menu, click Runtime and select **Change runtime type**. Set \"TPU\" as the hardware accelerator.\n",
        "* The cell below makes sure you have access to a TPU on Colab.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3P6b3uqfzpDI"
      },
      "source": [
        "import os\n",
        "assert os.environ['COLAB_TPU_ADDR'], 'Make sure to select TPU from Edit > Notebook settings > Hardware accelerator'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CHzziBW5AoZH"
      },
      "source": [
        "## Installing PyTorch/XLA\n",
        "\n",
        "Run the following cell (or copy it into your own notebook!) to install PyTorch, Torchvision, and PyTorch/XLA. It will take a couple minutes to run.\n",
        "\n",
        "The PyTorch/XLA package lets PyTorch connect to Cloud TPUs. (It's named PyTorch/XLA, not PyTorch/TPU, because XLA is the name of the TPU compiler.) In particular, PyTorch/XLA makes TPU cores available as PyTorch devices. This lets PyTorch create and manipulate tensors on TPUs."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yUB12htcqU9W"
      },
      "source": [
        "!pip install cloud-tpu-client==0.10 https://storage.googleapis.com/tpu-pytorch/wheels/torch_xla-1.7-cp36-cp36m-linux_x86_64.whl"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "__QtfkH8qVXC"
      },
      "source": [
        "### Only run the below commented cell if you would like a nightly release"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sPJVqAKyml5W"
      },
      "source": [
        "# VERSION = \"20200325\"  #@param [\"1.5\" , \"20200325\", \"nightly\"]\n",
        "# !curl https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o pytorch-xla-env-setup.py\n",
        "# !python pytorch-xla-env-setup.py --version $VERSION"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ls3j-EWI2D2v"
      },
      "source": [
        "## Creating and Manipulating Tensors on TPUs\n",
        "\n",
        "PyTorch uses Cloud TPUs just like it uses CPU or CUDA devices, as the next few cells will show. Each core of a Cloud TPU is treated as a different PyTorch  device.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "42avAvSg17by"
      },
      "source": [
        "# imports pytorch\n",
        "import torch\n",
        "\n",
        "# imports the torch_xla package\n",
        "import torch_xla\n",
        "import torch_xla.core.xla_model as xm"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4b8RfPPk4VIX"
      },
      "source": [
        "As mentioned above, the PyTorch/XLA package (torch_xla) lets PyTorch use TPU devices. The `xla_device()` function returns the TPU's \"default\" core as a device. This lets PyTorch creates tensors on TPUs:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "H9KYz-Vk4fMa"
      },
      "source": [
        "# Creates a random tensor on xla:1 (a Cloud TPU core)\n",
        "dev = xm.xla_device()\n",
        "t1 = torch.ones(3, 3, device = dev)\n",
        "print(t1)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yHwOC-xr4_LX"
      },
      "source": [
        "See the documentation at http://pytorch.org/xla/ for a description of all public PyTorch/XLA functions. Here `xm.xla_device()` acquired the first Cloud TPU core ('xla:1'). Other cores can be directly acquired, too:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2ef7flq95OxD"
      },
      "source": [
        "# Creating a tensor on the second Cloud TPU core\n",
        "second_dev = xm.xla_device(n=2, devkind='TPU')\n",
        "t2 = torch.zeros(3, 3, device = second_dev)\n",
        "print(t2)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ANcDKzGG5_ua"
      },
      "source": [
        "It is recommended that you use functions like `xm.xla_device()` over directly specifying TPU cores. A future Colab tutorial will show how to easily train a network using multiple cores (or you can look at [an example](https://colab.research.google.com/github/pytorch/xla/blob/master/contrib/colab/mnist-training.ipynb#scrollTo=Afwo4H7kSd8Phttps://)).\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QiKmkAkoO06x"
      },
      "source": [
        "Tensors on TPUs can be manipulated like any other PyTorch tensor. The following cell adds, multiplies, and matrix multiplies two tensors on a TPU core:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "l50-R2kwFY7Z"
      },
      "source": [
        "a = torch.randn(2, 2, device = dev)\n",
        "b = torch.randn(2, 2, device = dev)\n",
        "print(a + b)\n",
        "print(b * 2)\n",
        "print(torch.matmul(a, b))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cfDBbDtuisdu"
      },
      "source": [
        "This next cell runs a 1D convolution on a TPU core:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "aryiLyezisEg"
      },
      "source": [
        "# Creates random filters and inputs to a 1D convolution\n",
        "filters = torch.randn(33, 16, 3, device = dev)\n",
        "inputs = torch.randn(20, 16, 50, device = dev)\n",
        "torch.nn.functional.conv1d(inputs, filters)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "D5b4AmYDgbKd"
      },
      "source": [
        "And tensors can be transferred between CPU and TPU. In the following cell, a tensor on the CPU is copied to a TPU core, and then copied back to the CPU again. Note that PyTorch makes copies of tensors when transferring them across devices, so `t_cpu` and `t_cpu_again` are different tensors.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8WTsdQ3yO-8G"
      },
      "source": [
        "# Creates a tensor on the CPU (device='cpu' is unnecessary and only added for clarity)\n",
        "t_cpu = torch.randn(2, 2, device='cpu')\n",
        "print(t_cpu)\n",
        "\n",
        "t_tpu = t_cpu.to(dev)\n",
        "print(t_tpu)\n",
        "\n",
        "t_cpu_again = t_tpu.to('cpu')\n",
        "print(t_cpu_again)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DWtOgDLxN_BV"
      },
      "source": [
        "## Running PyTorch modules and autograd on TPUs\n",
        "\n",
        "Modules and autograd are fundamental PyTorch components. \n",
        "\n",
        "In PyTorch, every stateful function is a module. Modules are Python classes augmented with metadata that lets PyTorch understand how to use them in a neural network. For example, linear layers are modules, as are entire networks. Since modules are stateful, they can be placed on devices, too. PyTorch/XLA lets us place them on TPU cores:\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "P-WT-r8sRERM"
      },
      "source": [
        "# Creates a linear module\n",
        "fc = torch.nn.Linear(5, 2, bias=True)\n",
        "\n",
        "# Copies the module to the XLA device (the first Cloud TPU core)\n",
        "fc = fc.to(dev)\n",
        "\n",
        "# Creates a random feature tensor\n",
        "features = torch.randn(3, 5, device=dev, requires_grad=True)\n",
        "\n",
        "# Runs and prints the module\n",
        "output = fc(features)\n",
        "print(output)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "23Epn0HHR_Nq"
      },
      "source": [
        "Autograd is the system PyTorch uses to populate the gradients of weights in a neural network. See [here](https://pytorch.org/tutorials/beginner/blitz/autograd_tutorial.html#sphx-glr-beginner-blitz-autograd-tutorial-py) for details about PyTorch's autograd. When a module is run on a TPU core, its gradients are also populated on the same TPU core by autograd. The following cell demonstrates this:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Qs-2q5AMRixo"
      },
      "source": [
        "output.backward(torch.ones_like(output))\n",
        "print(fc.weight.grad)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "u0g3o1wHmF38"
      },
      "source": [
        "## Running PyTorch networks on TPUs\n",
        "\n",
        "As mentioned above, PyTorch networks are also modules, and so they're run in the same way. The following cell runs a relatively simple PyTorch network from the [PyTorch tutorial docs](https://pytorch.org/tutorials/beginner/blitz/neural_networks_tutorial.html#sphx-glr-beginner-blitz-neural-networks-tutorial-py) on a TPU core:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZLtM_M1imkFQ"
      },
      "source": [
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "\n",
        "# Simple example network from \n",
        "# https://pytorch.org/tutorials/beginner/blitz/neural_networks_tutorial.html#sphx-glr-beginner-blitz-neural-networks-tutorial-py\n",
        "class Net(nn.Module):\n",
        "\n",
        "    def __init__(self):\n",
        "        super(Net, self).__init__()\n",
        "        # 1 input image channel, 6 output channels, 3x3 square convolution\n",
        "        # kernel\n",
        "        self.conv1 = nn.Conv2d(1, 6, 3)\n",
        "        self.conv2 = nn.Conv2d(6, 16, 3)\n",
        "        # an affine operation: y = Wx + b\n",
        "        self.fc1 = nn.Linear(16 * 6 * 6, 120)  # 6*6 from image dimension\n",
        "        self.fc2 = nn.Linear(120, 84)\n",
        "        self.fc3 = nn.Linear(84, 10)\n",
        "\n",
        "    def forward(self, x):\n",
        "        # Max pooling over a (2, 2) window\n",
        "        x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))\n",
        "        # If the size is a square you can only specify a single number\n",
        "        x = F.max_pool2d(F.relu(self.conv2(x)), 2)\n",
        "        x = x.view(-1, self.num_flat_features(x))\n",
        "        x = F.relu(self.fc1(x))\n",
        "        x = F.relu(self.fc2(x))\n",
        "        x = self.fc3(x)\n",
        "        return x\n",
        "\n",
        "    def num_flat_features(self, x):\n",
        "        size = x.size()[1:]  # all dimensions except the batch dimension\n",
        "        num_features = 1\n",
        "        for s in size:\n",
        "            num_features *= s\n",
        "        return num_features\n",
        "\n",
        "\n",
        "# Places network on the default TPU core\n",
        "net = Net().to(dev)\n",
        "\n",
        "# Creates random input on the default TPU core\n",
        "input = torch.randn(1, 1, 32, 32, device=dev)\n",
        "\n",
        "# Runs network\n",
        "out = net(input)\n",
        "print(out)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zmySVJYIm88W"
      },
      "source": [
        "As in the previous snippets, running PyTorch on a TPU just requires specifying a TPU core as a device."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e9y2QGhl8SyE"
      },
      "source": [
        "## More PyTorch on TPUs!\n",
        "\n",
        "This \"Getting Started\" notebook showed you how to:\n",
        "\n",
        "* Install PyTorch/XLA on Colab, which lets you use PyTorch with TPUs.\n",
        "* Run basic PyTorch functions on TPUs, like creating and adding tensors.\n",
        "* Run PyTorch modules and autograd on TPUs.\n",
        "* Run PyTorch networks on TPUs.\n",
        "\n",
        "Working with tensors, running modules, and running entire networks on a Cloud TPU is as simple as installing PyTorch/XLA and telling PyTorch to use the Colab TPU as its device.\n",
        "\n",
        "You're encouraged to try PyTorch/XLA on Colab and GCP, too! You can copy this notebook to get started, or just copy the PyTorch/XLA setup cell above. The PyTorch/XLA team would love to hear your ideas/suggestions/comments, too! For now, the best way to contact the team is with an issue on our Github: https://github.com/pytorch/xla/issues. \n",
        "\n",
        "Future tutorials will cover:\n",
        "\n",
        "* Training a network using multiple TPU cores.\n",
        "* Loading batched data onto TPUs.\n",
        "\n",
        "You can already see examples of training a network on multiple TPU cores and loading batched data onto TPUs [here](https://github.com/pytorch/xla/tree/master/contrib/colab).\n",
        "\n"
      ]
    }
  ]
}