ray-train/03-Ray-Train-with-PyTorch.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Ray Train - A Library for Distributed Deep Learning\n",
    "\n",
    "© 2019-2021, Anyscale. All Rights Reserved\n",
    "\n",
    "![Anyscale Academy](../images/AnyscaleAcademyLogo.png)\n",
    "\n",
    "[Ray Train](https://docs.ray.io/en/latest/train/train.html) is a lightweight library for distributed deep learning. It provides thin wrappers around [PyTorch](https://pytorch.org) and [TensorFlow](https://tensorflow.org) native modules for data parallel training.\n",
    "\n",
    "> **NOTE**: Ray SGD is renamed to Ray Train"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## About Ray Train\n",
    "\n",
    "The main features of Ray Train are:\n",
    " * **Ease of use:** You can scale PyTorch’s native `DistributedDataParallel` and TensorFlow’s `tf.distribute.MirroredStrategy` without the requirement to monitor individual nodes yourself.\n",
    " * **Composability:** Ray Train is built on top of the [Ray Actor](https://docs.ray.io/en/latest/actors.html) API, enabling seamless integration with existing Ray applications such as RLlib, Tune, and Serve.\n",
    " * **Scale up and down:** You can start on a single CPU, then scale up to multi-node, multi-CPU, or multi-GPU clusters when needed. All it takes is changing two lines of code.\n",
    "\n",
    "This [Ray blog post](https://medium.com/distributed-computing-with-ray/faster-and-cheaper-pytorch-with-raysgd-a5a44d4fd220) provides more information on the motivations for Ray Train (SGD), such as the many steps you have to do yourself without it, and how it removes those steps."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example - Distributed Training for PyTorch \n",
    "\n",
    "This example is adapted and modified from the [Ray Train documentation](https://docs.ray.io/en/latest/train/examples/train_linear_example.html). \n",
    "\n",
    "First, do the necessary imports, as before."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import tqdm\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import ray.train as train\n",
    "from ray.train import Trainer\n",
    "from ray.train.torch import TorchConfig\n",
    "from ray.train.callbacks import JsonLoggerCallback, TBXLoggerCallback"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 1: Define PyTorch Datasets loaders "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now define classes and several functions we'll need.\n",
    "This particular tutorial uses linear regression to solve: `y = ax + b`.\n",
    "\n",
    "We implement our `LinearDataset` class as a PyTorch Dataset by subclassing `torch.utils.data.Dataset` "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "class LinearDataset(torch.utils.data.Dataset):\n",
    "    \"\"\"y = a * x + b\"\"\"\n",
    "\n",
    "    def __init__(self, a, b, size=1000):\n",
    "        x = np.arange(0, 10, 10 / size, dtype=np.float32)\n",
    "        self.X = torch.from_numpy(x)\n",
    "        self.y = torch.from_numpy(a * x + b)\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        return self.X[index, None], self.y[index, None]\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.X)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 2:  Define training and validation function per epoch \n",
    "\n",
    "Define our training function per epoch. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_epoch(dataloader, model, loss_fn, optimizer):\n",
    "    for X, y in dataloader:\n",
    "        # Compute prediction error\n",
    "        pred = model(X)\n",
    "        loss = loss_fn(pred, y)\n",
    "\n",
    "        # Backpropagation\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define our validate function per epoch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def validate_epoch(dataloader, model, loss_fn):\n",
    "    num_batches = len(dataloader)\n",
    "    model.eval()\n",
    "    loss = 0\n",
    "    with torch.no_grad():\n",
    "        for X, y in dataloader:\n",
    "            pred = model(X)\n",
    "            loss += loss_fn(pred, y).item()\n",
    "    loss /= num_batches\n",
    "    import copy\n",
    "    model_copy = copy.deepcopy(model)\n",
    "    result = {\"model\": model_copy.cpu().state_dict(), \"loss\": loss}\n",
    "    return result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 3: Define our training function to pass to the Trainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_func(config):\n",
    "    # Fetch all the configs\n",
    "    data_size = config.get(\"data_size\", 1000)\n",
    "    val_size = config.get(\"val_size\", 400)\n",
    "    batch_size = config.get(\"batch_size\", 32)\n",
    "    hidden_size = config.get(\"hidden_size\", 1)\n",
    "    lr = config.get(\"lr\", 1e-2)\n",
    "    epochs = config.get(\"epochs\", [20, 40, 60])\n",
    "\n",
    "    # Get the Training and validation dataset \n",
    "    train_dataset = LinearDataset(2, 5, size=data_size)\n",
    "    val_dataset = LinearDataset(2, 5, size=val_size)\n",
    "    \n",
    "    # Convert them to PyTorch equivalent dataloaders\n",
    "    # Prepare them to use for Ray Training distributed training\n",
    "    # by using the train.torch.prepare_data_loaders.\n",
    "    # These are wrappers around PyTorch Distributed Dataloaders\n",
    "    train_loader = torch.utils.data.DataLoader(\n",
    "        train_dataset, batch_size=batch_size)\n",
    "    validation_loader = torch.utils.data.DataLoader(\n",
    "        val_dataset, batch_size=batch_size)\n",
    "\n",
    "    train_loader = train.torch.prepare_data_loader(train_loader)\n",
    "    validation_loader = train.torch.prepare_data_loader(validation_loader)\n",
    "\n",
    "    # Create our simple PyTorch linear model and prepare it for PyTorch DDP\n",
    "    # \n",
    "    model = nn.Linear(1, hidden_size)\n",
    "    model = train.torch.prepare_model(model)\n",
    "\n",
    "    # Use MSE for our loss function\n",
    "    loss_fn = nn.MSELoss()\n",
    "\n",
    "    # Use SGD for optimzer\n",
    "    optimizer = torch.optim.SGD(model.parameters(), lr=lr)\n",
    "\n",
    "    results = []\n",
    "\n",
    "    for epoch in epochs:\n",
    "         for e in range(epoch):\n",
    "            train_epoch(train_loader, model, loss_fn, optimizer)\n",
    "            result = validate_epoch(validation_loader, model, loss_fn)\n",
    "\n",
    "            # Ray Train, as in Ray Tune, allows us to report results back\n",
    "            # to main Trainer\n",
    "            train.report(**result)\n",
    "            results.append(result)\n",
    "            if e % epoch == 0:\n",
    "                od = result.get('model')        #is an ordered dictionary\n",
    "                loss = result.get('loss')\n",
    "                m_weight = od.get('module.weight').item()\n",
    "                m_bias = od.get('module.bias').item()\n",
    "                \n",
    "                print(f\"epoch {epoch}, loss: {loss:.3f}, model.weight: {m_weight:.3f}, model.bias: {m_bias:.3f}\")\n",
    "\n",
    "    return results"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 4: Wrap our Trainer around a main driver function\n",
    "\n",
    "Also, note that we are using a PyTorch backend and providing a [TorchConfig](https://docs.ray.io/en/latest/train/api.html?highlight=TorchConfig#torchconfighttps://docs.ray.io/en/latest/train/api.html?highlight=TorchConfig#torchconfig) with [gloo](https://pytorch.org/docs/stable/distributed.htmlhttps://pytorch.org/docs/stable/distributed.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_linear(num_workers=2, use_gpu=False, epochs=[20, 40]):\n",
    "    trainer = Trainer(\n",
    "        backend=TorchConfig(backend=\"gloo\"),\n",
    "        num_workers=num_workers,\n",
    "        use_gpu=use_gpu)\n",
    "    config = {\"lr\": 1e-2, \"hidden_size\": 1, \"batch_size\": 4, \"epochs\": epochs}\n",
    "    trainer.start()\n",
    "    results = trainer.run(\n",
    "        train_func,\n",
    "        config,\n",
    "        callbacks=[JsonLoggerCallback(),\n",
    "                   TBXLoggerCallback()])\n",
    "    trainer.shutdown()\n",
    "    return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2022-03-16 16:14:37,498\tINFO services.py:1412 -- View the Ray dashboard at \u001b[1m\u001b[32mhttp://127.0.0.1:8267\u001b[39m\u001b[22m\n",
      "2022-03-16 16:14:39,998\tINFO trainer.py:199 -- Trainer logs will be logged in: /Users/jules/ray_results/train_2022-03-16_16-14-39\n",
      "2022-03-16 16:14:41,731\tINFO trainer.py:205 -- Run results will be logged in: /Users/jules/ray_results/train_2022-03-16_16-14-39/run_001\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=60974)\u001b[0m 2022-03-16 16:14:41,698\tINFO torch.py:66 -- Setting up process group for: env:// [rank=0, world_size=4]\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=60974)\u001b[0m [W ProcessGroupGloo.cpp:707] Warning: Unable to resolve hostname to a (local) address. Using the loopback address as fallback. Manually set the network interface to bind to with GLOO_SOCKET_IFNAME. (function operator())\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=60975)\u001b[0m 2022-03-16 16:14:41,698\tINFO torch.py:66 -- Setting up process group for: env:// [rank=3, world_size=4]\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=60975)\u001b[0m [W ProcessGroupGloo.cpp:707] Warning: Unable to resolve hostname to a (local) address. Using the loopback address as fallback. Manually set the network interface to bind to with GLOO_SOCKET_IFNAME. (function operator())\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=60980)\u001b[0m 2022-03-16 16:14:41,699\tINFO torch.py:66 -- Setting up process group for: env:// [rank=2, world_size=4]\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=60980)\u001b[0m [W ProcessGroupGloo.cpp:707] Warning: Unable to resolve hostname to a (local) address. Using the loopback address as fallback. Manually set the network interface to bind to with GLOO_SOCKET_IFNAME. (function operator())\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=60981)\u001b[0m 2022-03-16 16:14:41,698\tINFO torch.py:66 -- Setting up process group for: env:// [rank=1, world_size=4]\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=60981)\u001b[0m [W ProcessGroupGloo.cpp:707] Warning: Unable to resolve hostname to a (local) address. Using the loopback address as fallback. Manually set the network interface to bind to with GLOO_SOCKET_IFNAME. (function operator())\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=60974)\u001b[0m 2022-03-16 16:14:41,890\tINFO torch.py:244 -- Moving model to device: cpu\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=60974)\u001b[0m 2022-03-16 16:14:41,890\tINFO torch.py:247 -- Wrapping provided model in DDP.\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=60975)\u001b[0m 2022-03-16 16:14:41,890\tINFO torch.py:244 -- Moving model to device: cpu\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=60975)\u001b[0m 2022-03-16 16:14:41,890\tINFO torch.py:247 -- Wrapping provided model in DDP.\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=60980)\u001b[0m 2022-03-16 16:14:41,890\tINFO torch.py:244 -- Moving model to device: cpu\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=60980)\u001b[0m 2022-03-16 16:14:41,890\tINFO torch.py:247 -- Wrapping provided model in DDP.\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=60981)\u001b[0m 2022-03-16 16:14:41,890\tINFO torch.py:244 -- Moving model to device: cpu\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=60981)\u001b[0m 2022-03-16 16:14:41,890\tINFO torch.py:247 -- Wrapping provided model in DDP.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=60974)\u001b[0m epoch 20, loss: 3.587, model.weight: 2.326, model.bias: 1.743\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=60975)\u001b[0m epoch 20, loss: 3.507, model.weight: 2.326, model.bias: 1.743\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=60981)\u001b[0m epoch 20, loss: 3.560, model.weight: 2.326, model.bias: 1.743\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=60980)\u001b[0m epoch 20, loss: 3.534, model.weight: 2.326, model.bias: 1.743\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=60974)\u001b[0m epoch 40, loss: 0.004, model.weight: 2.011, model.bias: 4.886\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=60975)\u001b[0m epoch 40, loss: 0.004, model.weight: 2.011, model.bias: 4.886\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=60981)\u001b[0m epoch 40, loss: 0.004, model.weight: 2.011, model.bias: 4.886\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=60980)\u001b[0m epoch 40, loss: 0.004, model.weight: 2.011, model.bias: 4.886\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=60974)\u001b[0m epoch 60, loss: 0.000, model.weight: 2.000, model.bias: 5.000\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=60975)\u001b[0m epoch 60, loss: 0.000, model.weight: 2.000, model.bias: 5.000\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=60981)\u001b[0m epoch 60, loss: 0.000, model.weight: 2.000, model.bias: 5.000\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=60980)\u001b[0m epoch 60, loss: 0.000, model.weight: 2.000, model.bias: 5.000\n"
     ]
    }
   ],
   "source": [
    "results = train_linear(\n",
    "            num_workers=4,\n",
    "            epochs=[20, 40, 60])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Logging callbacks\n",
    "\n",
    "They store the results in `~/ray_results/train_*`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls -l ~/ray_results/train_*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Launch Tensorboard to view the training results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "NOTE: Using experimental fast data loading logic. To disable, pass\n",
      "    \"--load_fast=false\" and report issues on GitHub. More details:\n",
      "    https://github.com/tensorflow/tensorboard/issues/4784\n",
      "\n",
      "Serving TensorBoard on localhost; to expose to the network, use a proxy or pass --bind_all\n",
      "TensorBoard 2.8.0 at http://localhost:6006/ (Press CTRL+C to quit)\n",
      "^C\n"
     ]
    }
   ],
   "source": [
    "!tensorboard --logdir ~/ray_results/train_2021-12-12_16-20-46/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"images/ray_train_tensorboard.png\" width=\"80%\" height=\"80%\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Excercises\n",
    "\n",
    "Have a go at this in your spare time and observe the results\n",
    "\n",
    " 1. Change the learning rate and batch size in `config`\n",
    " 2. Try chaning the number of workers to 1/2 number of cores on your localhost or laptop\n",
    " 3. Change the `data_size` and `val_size`\n",
    " 4. Modify the linear equation: `y = 2x + 5 --> y = 4x + 10`. This will require you to modify `LinearDataset(2, 5, size=data_size)`"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "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.8.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}