pytorch_resnet_singlecore.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "82Yd9t9H9cjt"
   },
   "source": [
    "### Single-core Training FashionMNIST on Cloud TPU\n",
    "\n",
    "In this notebook we explore the simple code changes necessary to train and evaluate a model on a TPU.\n",
    "\n",
    "Some of the some of the torch_xla functions we will use include:\n",
    "* `device = xm.xla_device()` to set the device to TPU.\n",
    "* `xm.is_master_ordinal()` to make sure the current process is the master ordinal (0).\n",
    "* `xm.optimizer_step()` for evaluation each time gradients are updated."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/rastringer/TPU_examples/blob/main/pytorch_resnet_singlecore.ipynb)\n",
    "\n",
    "<table align=\"left\">\n",
    "  <td>\n",
    "    <a href=\"https://github.com/rastringer/rastringer/TPU_examples/blob/main/pytorch_resnet_singlecore.ipynb\">\n",
    "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
    "      View on GitHub\n",
    "    </a>\n",
    "  </td>\n",
    "  <td>\n",
    "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/rastringer/TPU_examples/blob/main/pytorch_resnet_singlecore.ipynb\">\n",
    "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
    "      Open in Vertex AI Workbench\n",
    "    </a>\n",
    "  </td>\n",
    "</table>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you are running this notebook on **Vertex Workbench** or in another **Jupyter** environment, proceed to the **Training** subhead and skip the following cells for setting up on Colab. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "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",
    "* Uncomment the cell below to make sure you have access to a TPU on Colab."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import os\n",
    "# assert os.environ['COLAB_TPU_ADDR'], 'Make sure to select TPU from Edit > Notebook settings > Hardware accelerator'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Install Colab TPU compatible PyTorch/TPU wheels and dependencies\n",
    "\n",
    "Uncomment and run if using Colab."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install cloud-tpu-client==0.10 torch==2.0.0 torchvision==0.15.1 https://storage.googleapis.com/tpu-pytorch/wheels/colab/torch_xla-2.0-cp39-cp39-linux_x86_64.whl"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch_xla\n",
    "import torch_xla.core.xla_model as xm\n",
    "import torch_xla.distributed.xla_multiprocessing as xmp\n",
    "import torch_xla.distributed.parallel_loader as pl\n",
    "import torchvision\n",
    "from torchvision import datasets\n",
    "import torchvision.transforms as transforms\n",
    "import time\n",
    "import sys"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Training parameters\n",
    "flags = {}\n",
    "flags['batch_size'] = 8\n",
    "flags['num_epochs'] = 1\n",
    "flags['seed'] = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Random seed for initialization\n",
    "torch.manual_seed(flags['seed'])\n",
    "# Sets device to Cloud TPU core\n",
    "device = xm.xla_device()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_datasets():\n",
    "    \n",
    "    # Normalization for dataloader \n",
    "    # TorchVision models require RGB (3 x H x W) images\n",
    "    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],\n",
    "                                  std=[0.229, 0.224, 0.225])\n",
    "    to_rgb = transforms.Lambda(lambda image: image.convert('RGB'))\n",
    "    resize = transforms.Resize((224, 224))\n",
    "    my_transform = transforms.Compose([resize, to_rgb, transforms.ToTensor(), normalize])\n",
    "\n",
    "    # Checks if current process is the master ordinal (0)\n",
    "    # Other workers wait for master to complete download\n",
    "    if not xm.is_master_ordinal():\n",
    "        xm.rendezvous('download_only_once')\n",
    "\n",
    "    train_dataset = datasets.FashionMNIST(\n",
    "        \"/tmp/fashionmnist\",\n",
    "        train=True,\n",
    "        download=True,\n",
    "        transform=my_transform)\n",
    "\n",
    "    test_dataset = datasets.FashionMNIST(\n",
    "        \"/tmp/fashionmnist\",\n",
    "        train=False,\n",
    "        download=True,\n",
    "        transform=my_transform)\n",
    "\n",
    "    return train_dataset, test_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_datasets():\n",
    "    \n",
    "    train_dataset, test_dataset = create_datasets()\n",
    "    \n",
    "    train_sampler = torch.utils.data.RandomSampler(train_dataset)\n",
    "    test_sampler = torch.utils.data.RandomSampler(test_dataset)\n",
    "\n",
    "    # Dataloaders load data in batches\n",
    "    train_loader = torch.utils.data.DataLoader(\n",
    "        train_dataset,\n",
    "        batch_size=flags['batch_size'],\n",
    "        sampler=train_sampler\n",
    "    )\n",
    "\n",
    "    test_loader = torch.utils.data.DataLoader(\n",
    "        test_dataset,\n",
    "        batch_size=flags['batch_size'],\n",
    "        sampler=test_sampler\n",
    "    )\n",
    "    \n",
    "    return train_loader, test_loader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Model, optimizer, and loss function \n",
    "# We use a resnet18 model for the 10 classes of the FashionMNIST dataset\n",
    "net = torchvision.models.resnet18(num_classes=10).to(device).train()\n",
    "loss_fn = torch.nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.Adam(net.parameters())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train():\n",
    "    train_start = time.time()\n",
    "    train_loader, test_loader = load_datasets()\n",
    "    for data, targets in iter(train_loader):\n",
    "        # Sends data and targets to device\n",
    "        data = data.to(device)\n",
    "        targets = targets.to(device)\n",
    "        # Get prediction\n",
    "        output = net(data)\n",
    "        # Loss function\n",
    "        loss = loss_fn(output, targets)\n",
    "        # Update model\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        # The 'barrier' here forces evaluation each time gradients are \n",
    "        # updates, keeping graphs manageable.\n",
    "        xm.optimizer_step(optimizer, barrier=True)\n",
    "\n",
    "    elapsed_train_time = time.time() - train_start\n",
    "    print(f\"Finished training. Train time was: {elapsed_train_time}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test():\n",
    "    net.eval()\n",
    "    eval_start = time.time()\n",
    "    with torch.no_grad():\n",
    "        num_correct = 0\n",
    "        total_guesses = 0\n",
    "    train_loader, test_loader = load_datasets()\n",
    "    for data, targets in iter(test_loader):\n",
    "        # Sends data and targets to device\n",
    "        data = data.to(device)\n",
    "        targets = targets.to(device)\n",
    "        output = net(data)\n",
    "        best_guesses = torch.argmax(output, 1)\n",
    "        # Calculate accuracy\n",
    "        num_correct += torch.eq(targets, best_guesses).sum().item()\n",
    "        total_guesses += flags['batch_size']\n",
    "\n",
    "    accuracy = 100.0 * num_correct / total_guesses\n",
    "    elapsed_eval_time = time.time() - eval_start\n",
    "    print(f\"Finished evaluation. Evaluation time was: {elapsed_eval_time}\")\n",
    "    print(f\"Guessed {num_correct} of {total_guesses} correctly for {accuracy} % accuracy.\")\n",
    "\n",
    "    return accuracy, data, targets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def trainer():\n",
    "\n",
    "    accuracy = 0.0\n",
    "    data, targets = None, None\n",
    "\n",
    "    # Loop through epochs, calling the train and eval functions above\n",
    "    for epoch in range(flags['num_epochs']):\n",
    "        tic = time.time()\n",
    "        train()\n",
    "        xm.master_print(\"Finished training epoch {}\".format(epoch))\n",
    "        toc = time.time()\n",
    "        accuracy, data, targets  = test()\n",
    "        # Calculate training time per epoch\n",
    "        epoch_time = round(toc-tic, 2)\n",
    "        print(f\"Epoch trained on a single TPU in {epoch_time} seconds\")\n",
    "        \n",
    "    return accuracy, data, targets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this notebook, we trained and evaluated a ResNet18 model on the FashionMNIST dataset.\n",
    "We used various functions from the pytorch_XLA library to run on a Cloud TPU.\n",
    "\n",
    "Try experimenting with the `batch_size` - how high can it go before the training starts to slow down?\n",
    "\n",
    "We can also experiment with different model architectures, such as EfficientNet or ConvNeXt, just by changing this line:\n",
    "\n",
    "```\n",
    "net = torchvision.models.resnet18(num_classes=10).to(device).train()\n",
    "\n",
    "```\n",
    "\n",
    "to swap `resnet18` to others listed [here](https://pytorch.org/vision/stable/models.html#classification).\n",
    "                                           \n",
    "The next notebook, `pytorch_resnet_multicore.ipynb`, shows how we can adapt our code to run on multiple TPU cores, which should\n",
    "make training more efficient."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To save this notebook as a Python file, run the following commmand from a terminal in the same folder:\n",
    "```\n",
    "jupyter nbconvert --to script pytorch_resnet_singlecore.ipynb\n",
    "```\n",
    "\n",
    "We can then load the .py file into a Cloud Storage bucket, then import it via `gsutil` into the TPU VM.\n",
    "```\n",
    "gsutil cp pytorch_resnet_singlecore.py gs://<your-unique-bucket-name>\n",
    "```\n",
    "\n",
    "Then write this command on the TPU VM terminal:\n",
    "\n",
    "```\n",
    "gsutil cp gs://<your-unique-bucket-name>/pytorch_resnet_singlecore.py .\n",
    "\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "TPU",
  "colab": {
   "provenance": []
  },
  "gpuClass": "standard",
  "kernelspec": {
   "display_name": "Pytorch (Local)",
   "language": "python",
   "name": "local-pytorch"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}