pytorch_resnet_multicore.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "82Yd9t9H9cjt"
   },
   "source": [
    "### Multi-core Training FashionMNIST on Cloud TPU\n",
    "\n",
    "Now we are familiar with adjusting code to run on a single TPU core, we will walkthrough expanding the data loading, model training and evaluation across all eight cores of the TPU.\n",
    "Multi-core operations will predictibly be faster and able to handle larger loads (datasets, batch sizes).\n",
    "\n",
    "Beyond setting `device = xm.xla_device()`, below are some of the torch_xla functions we will use:\n",
    "* Data is only being downloaded once by a master worker by checking `xm.is_master_ordinal()`.\n",
    "* Subsets of the data are being loaded efficiently across all processes using `DistributedSampler`.\n",
    "* `xm.optimizer_step(optimizer)` to consolidates gradients between cores during training.\n",
    "* We run the training using `xmp.spawn()` to enable replication across multiple processes."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "OuiXOhDI8lrp"
   },
   "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_multicore.ipynb)\n",
    "\n",
    "<table align=\"left\">\n",
    "  <td>\n",
    "    <a href=\"https://github.com/rastringer/rastringer/TPU_examples/blob/main/pytorch_resnet_multicore.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_multicore.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": {
    "id": "tOzBiGh38lrp"
   },
   "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": {
    "id": "k3XRr3hT8lrq"
   },
   "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": {
    "id": "WEElDtuP8lrq"
   },
   "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": {
    "id": "ToGe-TPy8lrr"
   },
   "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": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "yJtFtdsx8lrs",
    "outputId": "85a7e835-464b-48af-a274-23d88ecf0d79"
   },
   "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": [
    "### Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "zH08bH2F_gaq"
   },
   "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": "markdown",
   "metadata": {},
   "source": [
    "### Training flags"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "WdiFzm2T_gdY"
   },
   "outputs": [],
   "source": [
    "# Training parameters\n",
    "flags = {}\n",
    "flags['batch_size'] = 32\n",
    "flags['num_workers'] = 8\n",
    "flags['num_epochs'] = 1\n",
    "flags['seed'] = 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Device and seed initialization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "CWdfwO_e_gfq"
   },
   "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": "markdown",
   "metadata": {
    "id": "ggUz7tW28lrs"
   },
   "source": [
    "### Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "PPHQIvhT_5JK"
   },
   "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\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "xNxJnJ23Aqzo",
    "outputId": "f6986f00-3d20-44fa-a884-bbd1476ff353"
   },
   "outputs": [],
   "source": [
    "train_dataset, test_dataset = create_datasets()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "q_Fjwl6iAKii"
   },
   "outputs": [],
   "source": [
    "def load_datasets():\n",
    "    \n",
    "  # Ensure data is downloaded only once by a master worker   \n",
    "  if xm.is_master_ordinal():\n",
    "        xm.rendezvous('download_only_once')\n",
    "\n",
    "  # DistributedSampler restricts data loading to a subset of the dataset\n",
    "  # for each process\n",
    "  train_sampler = torch.utils.data.distributed.DistributedSampler(\n",
    "      train_dataset,\n",
    "      num_replicas=xm.xrt_world_size(),\n",
    "      rank=xm.get_ordinal(),\n",
    "      shuffle=True)\n",
    "\n",
    "  test_sampler = torch.utils.data.distributed.DistributedSampler(\n",
    "      test_dataset,\n",
    "      num_replicas=xm.xrt_world_size(),\n",
    "      rank=xm.get_ordinal(),\n",
    "      shuffle=False)\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",
    "      num_workers=flags['num_workers'],\n",
    "      drop_last=True)\n",
    "\n",
    "  test_loader = torch.utils.data.DataLoader(\n",
    "      test_dataset,\n",
    "      batch_size=flags['batch_size'],\n",
    "      sampler=test_sampler,\n",
    "      shuffle=False,\n",
    "      num_workers=flags['num_workers'],\n",
    "      drop_last=True)\n",
    "\n",
    "  return train_loader, test_loader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "lMxckRVWBi-g",
    "outputId": "2468189b-c307-4fb5-d5ec-5b78c3642469"
   },
   "outputs": [],
   "source": [
    "train_loader, test_loader = load_datasets()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model, optimizer, and loss function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "B68mZjHQAKk5"
   },
   "outputs": [],
   "source": [
    "# We use a resnet18 model for the 10 classes of the\n",
    "# FashionMNIST dataset\n",
    "# Each process has its own copy of the model\n",
    "net = torchvision.models.resnet18(num_classes=10).to(device).train()\n",
    "loss_fn = torch.nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.Adam(net.parameters())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "KAFOdTSOA3Rq"
   },
   "outputs": [],
   "source": [
    "def train(loader):\n",
    "    \n",
    "  for batch_num, batch in enumerate(loader):\n",
    "    data, targets = batch\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",
    "    # xm.optimizer_step(optimizer) consolidates the gradients between cores\n",
    "    # and issues the XLA device step computation.\n",
    "    xm.optimizer_step(optimizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "y447qhZwA3UP"
   },
   "outputs": [],
   "source": [
    "def test(loader):\n",
    "    net.eval()\n",
    "    eval_start = time.time()\n",
    "    with torch.no_grad():\n",
    "        num_correct = 0\n",
    "        total_guesses = 0\n",
    "\n",
    "    for batch_num, batch in enumerate(loader):\n",
    "        data, targets = batch\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": {
    "id": "sxpClYQN8lrt"
   },
   "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",
    "        epoch_start = time.time()\n",
    "        # ParallelLoader wraps a DataLoader with background data upload\n",
    "        para_train_loader = pl.ParallelLoader(train_loader, [device]).per_device_loader(device)\n",
    "        train(para_train_loader)\n",
    "        epoch_time = time.time() - epoch_start\n",
    "        # print(f\"Finished training epoch in {epoch_time}\")\n",
    "        xm.master_print(f\"Finished training epoch in {epoch_time} seconds \")\n",
    "        para_test_loader = pl.ParallelLoader(test_loader, [device]).per_device_loader(device)\n",
    "        accuracy, data, targets  = test(para_test_loader)\n",
    "    \n",
    "    return accuracy, data, targets\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "uoSQeQN28lru"
   },
   "source": [
    "The *multiprocess function* runs the trainer, takes the index of each process and the flags defined above"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "4j11tY5Y8lru",
    "outputId": "af549bff-9f47-44ba-8aec-80f80825faac"
   },
   "outputs": [],
   "source": [
    "def _mp_fn(index, flags):\n",
    "    torch.set_default_tensor_type('torch.FloatTensor')\n",
    "    accuracy, data, targets = trainer()\n",
    "\n",
    "# Enable replication across multiple processes\n",
    "xmp.spawn(_mp_fn, args=(flags,), nprocs=1, start_method='fork')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "K8k9mIye8lru"
   },
   "source": [
    "In this notebook, we trained and evaluated a ResNet18 model on the FashionMNIST dataset.\n",
    "We adjusted the code from the single-core example to include various data and model parallelization features that should have sped up our operations as they were distributed across 8 TPU cores. We can check the time against the single-core example, and predictably should see that the epochs train faster.\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)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DFcGPvYp8lrv"
   },
   "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_multicore.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",
    "Copy into Storage:\n",
    "\n",
    "```\n",
    "gsutil cp pytorch_resnet_multicore.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_multicore.py .\n",
    "\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "TEU8Rror8lrv"
   },
   "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
}