ray-train/01-Ray-Train-Quickstart.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "077f02f4-da74-4a16-8cb9-62509802ca3d",
   "metadata": {},
   "source": [
    "# Ray Train - A Library for Distributed Deep Learning\n",
    "\n",
    "© 2019-2022, 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), [TensorFlow](https://tensorflow.org), and [Horvod](https://horovod.ai/) native modules for data parallel training.\n",
    "\n",
    "> **NOTE**: Ray SGD is renamed to Ray Train"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9a43426-da48-4463-88b4-8a3dc9dd9843",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Quick Start: single machine, single worker with PyTorch\n",
    "\n",
    "Let's work through a typical non-distributed PyTorch trainining example, where we only use a single machine, single workper process."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "9f46a636-862c-47fd-8a62-4ce3b762bade",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dbff7b76-9026-4315-9786-22c24e41e430",
   "metadata": {},
   "source": [
    "### Step 1. Define constants, input and output variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "f17af003-04bb-47b6-bfbf-c7af92c53ce4",
   "metadata": {},
   "outputs": [],
   "source": [
    "NUM_SAMPLES = 20             # our dataset for training\n",
    "INPUT_SIZE = 20              # inputs or neurons into the first layer\n",
    "LAYER_SIZE = 15              # inputs or neurons to the hidden layer\n",
    "OUTPUT_SIZE = 5              # outputs to the last layer\n",
    "\n",
    "# In this example we use a randomly generated dataset.\n",
    "input = torch.randn(NUM_SAMPLES, INPUT_SIZE)         # In normal ML parlance, X\n",
    "labels = torch.randn(NUM_SAMPLES, OUTPUT_SIZE)       # In nmormal ML parlance, y"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39b1dfec-e469-4acd-9e95-70355b8450dc",
   "metadata": {},
   "source": [
    "### Step 2: Define a simple PyTorch neural network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "de5a74c7-c25c-490f-8d4b-a5bafbbcff6f",
   "metadata": {},
   "outputs": [],
   "source": [
    "class NeuralNetwork(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(NeuralNetwork, self).__init__()\n",
    "        self.layer1 = nn.Linear(in_features=INPUT_SIZE, out_features=LAYER_SIZE)\n",
    "        # Our activation function\n",
    "        self.relu = nn.ReLU()           \n",
    "        self.layer2 = nn.Linear(in_features=LAYER_SIZE, out_features=OUTPUT_SIZE)\n",
    "\n",
    "    def forward(self, input):\n",
    "        return self.layer2(self.relu(self.layer1(input)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82084e6c-03c8-4f71-bab9-3d7518e5c28f",
   "metadata": {},
   "source": [
    "### Step 3: Define your training function\n",
    "Simple function that iterates over epochs and does standard PyTorch steps:\n",
    " * Invoke the callable model with input\n",
    " * Calculate the loss\n",
    " * Zero out the gradients\n",
    " * Do backward propogation\n",
    " * Optimize the step"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "c7439b1f-3821-4916-851a-9eb883fab0dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_func(configs):\n",
    "    model = NeuralNetwork()\n",
    "    loss_fn = nn.MSELoss()\n",
    "    optimizer = optim.SGD(model.parameters(), lr=0.1)\n",
    "\n",
    "    # Iterate over the loop\n",
    "    epochs = configs.get('NUM_EPOCHS',[20, 40, 60] )\n",
    "    for epoch in epochs: \n",
    "        for e in tqdm(range(epoch)):\n",
    "            output = model(input)\n",
    "            loss = loss_fn(output, labels)\n",
    "            \n",
    "            optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "        \n",
    "            if e % epoch == 0:\n",
    "                print(f'epoch {epoch}, loss: {loss.item():.3f}')\n",
    "    # Return anything you want; here we just report back the pid of the Ray worker process on which this function runs\n",
    "    return os.getpid()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0811236e-a356-441f-ad52-2a81da4a5d61",
   "metadata": {},
   "source": [
    "### Step 4: Train the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "2f266353-fa21-48ae-a5db-85d3b0d07030",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 20/20 [00:00<00:00, 1870.54it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 20, loss: 0.923\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 40/40 [00:00<00:00, 2809.22it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 40, loss: 0.599\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 60/60 [00:00<00:00, 3168.26it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 60, loss: 0.277\n",
      "pid: 60743\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "result = train_func({'NUM_EPOCHS': [20, 40, 60]})\n",
    "print(f'pid: {result}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f90a70e3-245a-454c-a319-7203d0110ffe",
   "metadata": {},
   "source": [
    "### Excercises\n",
    "\n",
    "Have a a go at this in your spare time and observe the results\n",
    "\n",
    " 1. Change the NUM_EPOCHS list to **[200, 400, 600]**\n",
    " 2. Do you see the loss approaching zero?\n",
    " 3. Try changing sample sizes. Do you need more epochs to train and minimize loss?"
   ]
  }
 ],
 "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": 5
}