ray-train/04-Ray-Train-with-PyTorch-FashionMNIST.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9ef3ff2e-b5ec-466f-86d8-f08eb2f1402d",
   "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": "f4f3eda8-7185-48d4-9ff1-679621b987fd",
   "metadata": {},
   "source": [
    "## PyTorch Fashion MNIST for Distributed Training\n",
    "\n",
    "<img src=\"images/fashion-mnist-sprite.jpeg\" width=\"70%\" height=\"60%\"> \n",
    "\n",
    "We will use Ray Train to distribute our training using couple of models and evaluating which of the two provides us\n",
    "the best accuracy and a minimal loss. \n",
    "\n",
    "As excercise, you can try to further investigate how you improve the model—via regularization techniques, using CNN layers, trying different loss functions.\n",
    "\n",
    "The steps we will follow are no different (may be slight variation but the essence is the same) from the previous notenbooks.\n",
    "\n",
    "So let's go!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4755b099-304f-49dd-a935-303133d42126",
   "metadata": {},
   "source": [
    "First, do the necessary imports, as before."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f16002ec-9ec0-4c26-bee0-c839601aea35",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from typing import Dict\n",
    "\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "\n",
    "import ray\n",
    "import ray.train as train\n",
    "from ray.train.trainer import Trainer\n",
    "from ray.train.callbacks import JsonLoggerCallback, TBXLoggerCallback\n",
    "from torch import nn\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision import datasets\n",
    "from torchvision.transforms import ToTensor"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35e9e3ab-d950-42a0-adf2-cfa57eaec74d",
   "metadata": {},
   "source": [
    "### Step 1: Download Train and test datasets "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "a05d4628-08c2-4d4f-89aa-939eb819a437",
   "metadata": {},
   "outputs": [],
   "source": [
    "training_data = datasets.FashionMNIST(\n",
    "    root=\"~/data\",\n",
    "    train=True,\n",
    "    download=True,\n",
    "    transform=ToTensor(),\n",
    ")\n",
    "# Download test data from open datasets.\n",
    "test_data = datasets.FashionMNIST(\n",
    "    root=\"~/data\",\n",
    "    train=False,\n",
    "    download=True,\n",
    "    transform=ToTensor(),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd077dca-a562-4625-af4e-f0c5e068c712",
   "metadata": {},
   "source": [
    "## Step 2: Define a Neural Network Models. \n",
    "\n",
    "This is a quite simple NN model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "3f805792-c35a-4fef-b8bc-53fbea7eecd4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define model-1\n",
    "class NeuralNetwork(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(NeuralNetwork, self).__init__()\n",
    "        self.flatten = nn.Flatten()\n",
    "        self.linear_relu_stack = nn.Sequential(\n",
    "            nn.Linear(28 * 28, 512), nn.ReLU(), nn.Linear(512, 512), nn.ReLU(),\n",
    "            nn.Linear(512, 10), nn.ReLU())\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.flatten(x)\n",
    "        logits = self.linear_relu_stack(x)\n",
    "        return logits"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22d6ad67-8968-4c6d-a106-708e704d67b4",
   "metadata": {},
   "source": [
    "Define a deeper NN model archiecture with dropouts\n",
    "\n",
    "<img src=\"https://miro.medium.com/max/1400/1*2SHOuTUK51_Up3D9JMAplA.png\" width=\"70%\" height=\"50%\">\n",
    "\n",
    "[source](https://medium.com/@aaysbt/fashion-mnist-data-training-using-pytorch-7f6ad71e96f4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "653205a7-d6e3-4e5f-9627-42b69950529f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define model-2\n",
    "class Classifier(nn.Module):\n",
    "  def __init__(self):\n",
    "    super().__init__()\n",
    "    self.fc1 = nn.Linear(784, 120)\n",
    "    self.fc2 = nn.Linear(120, 120)\n",
    "    self.fc3 = nn.Linear(120,10)\n",
    "    self.dropout = nn.Dropout(0.2)\n",
    "\n",
    "  def forward(self,x):\n",
    "    x = x.view(x.shape[0],-1)\n",
    "    x = self.dropout(F.relu(self.fc1(x)))\n",
    "    x = self.dropout(F.relu(self.fc2(x)))\n",
    "    x = F.log_softmax(self.fc3(x), dim=1)\n",
    "    return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "bf7695d5-7ce3-44f3-a602-b9e04a471dff",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define accuracy function\n",
    "def accuracy_fn(y_pred, y_true):\n",
    "    n_correct = torch.eq(y_pred, y_true).sum().item()\n",
    "    acc = (n_correct / len(y_pred)) * 100\n",
    "    return acc"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37c6bc36-8344-4792-b9ee-c89c9f741016",
   "metadata": {},
   "source": [
    "### Step 3: Define per epoch training and validation functinos"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "cf218b87-a60c-4ee4-8952-7cf3a0abb5ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_epoch(dataloader, model, loss_fn, optimizer, epoch):\n",
    "    size = len(dataloader.dataset)\n",
    "    for batch, (X, y) in enumerate(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": "code",
   "execution_count": 7,
   "id": "314c237f-f16e-47f3-b2d1-17f9c6eaa075",
   "metadata": {},
   "outputs": [],
   "source": [
    "def validate_epoch(dataloader, model, loss_fn, epoch):\n",
    "    size = len(dataloader.dataset)\n",
    "    num_batches = len(dataloader)\n",
    "    model.eval()\n",
    "    test_loss, correct, acc =  0, 0, 0.0\n",
    "    with torch.no_grad():\n",
    "        for X, y in dataloader:\n",
    "            pred = model(X)\n",
    "            test_loss += loss_fn(pred, y).item()\n",
    "            correct += (pred.argmax(1) == y).type(torch.float).sum().item()\n",
    "            predictions = pred.max(dim=1)[1]\n",
    "            acc += accuracy_fn(predictions, y)\n",
    "    test_loss /= num_batches\n",
    "    acc /= num_batches\n",
    "    correct /= size\n",
    "    if epoch > 0 and epoch % 50 == 0:\n",
    "        print(f\"Epoc: {epoch}, Avg validation loss: {test_loss:.2f}, Avg validation accuracy: {acc:.2f}%\") \n",
    "        print(\"--\" * 40)\n",
    "    return test_loss"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1671c44a-7def-44b8-a468-704694c61d91",
   "metadata": {},
   "source": [
    "### Step 4: Define Ray Train Training function\n",
    "This function will be passed to `train.run(...)`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "8d2ccc74-555f-40c7-a54a-75506c7cbf25",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_func(config: Dict):\n",
    "    batch_size = config.get(\"batch_size\", 64) \n",
    "    lr = config.get('lr', 1e-3)\n",
    "    epochs = config.get(\"epochs\", 20)\n",
    "    momentum = config.get(\"momentum\", 0.9)\n",
    "    model_type = config.get('model_type', None)\n",
    "    loss_fn = config.get(\"loss_fn\", nn.NLLLoss())\n",
    "\n",
    "    # Create data loaders.\n",
    "    train_dataloader = DataLoader(training_data, batch_size=batch_size)\n",
    "    test_dataloader = DataLoader(test_data, batch_size=batch_size)\n",
    "\n",
    "    # Prepare to use Ray integrated wrappers around PyTorch's Dataloaders\n",
    "    train_dataloader = train.torch.prepare_data_loader(train_dataloader)\n",
    "    test_dataloader = train.torch.prepare_data_loader(test_dataloader)\n",
    "\n",
    "    # Create model.\n",
    "\n",
    "    model = Classifier() if model_type else NeuralNetwork()\n",
    "    # Prepare to use Ray integrated wrappers around PyTorch's model\n",
    "    model = train.torch.prepare_model(model)\n",
    "    \n",
    "    # Get or objective loss function\n",
    "    loss_fn = config.get(\"loss_fn\", nn.NLLLoss())\n",
    "\n",
    "    optimizer = torch.optim.SGD(model.parameters(), lr=lr, momentum=momentum)\n",
    "\n",
    "    loss_results = []\n",
    "\n",
    "    for e in range(epochs):\n",
    "        train_epoch(train_dataloader, model, loss_fn, optimizer, e)\n",
    "        loss = validate_epoch(test_dataloader, model, loss_fn, e)\n",
    "        train.report(loss=loss)\n",
    "        loss_results.append(loss)\n",
    "\n",
    "    return loss_results"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3242704c-0ec7-43e6-a7b8-27fa37dedb53",
   "metadata": {},
   "source": [
    "### Step 5: Wrap our Trainer around a main driver function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "9280b0c5-c754-4416-83ec-1f02c3028483",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_fashion_mnist(num_workers=12, use_gpu=False):\n",
    "    trainer = Trainer(\n",
    "        backend=\"torch\", num_workers=num_workers, use_gpu=use_gpu)\n",
    "    trainer.start()\n",
    "    result = trainer.run(\n",
    "        train_func=train_func,\n",
    "        config={\n",
    "            \"lr\": 1e-3,\n",
    "            \"batch_size\": 128,\n",
    "            \"epochs\": 150,\n",
    "            \"momentum\": 0.9,\n",
    "            \"model_type\": 0,                     # change to 1 for second NN model\n",
    "            \"loss_fn\": nn.CrossEntropyLoss()     # change to nn.nn.NLLLoss() \n",
    "        },\n",
    "        callbacks=[JsonLoggerCallback(), TBXLoggerCallback()])\n",
    "    trainer.shutdown() \n",
    "    return result"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6a6ab92-10c2-4594-84a9-f475ba634a0a",
   "metadata": {},
   "source": [
    "### Step 6: Define some parallelism parameters \n",
    "And a URL to connect to a Ray Cluster if running on Anysacle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "fc5b6b09-54f8-4485-b5e0-8a3051e5d45b",
   "metadata": {},
   "outputs": [],
   "source": [
    "number_of_workers = 8\n",
    "use_gpu = False                              # change to True if using a Ray cluster with GPUs\n",
    "address = \"anyscale://ray_train_ddp_cluster\" # use your anyscale cluster here"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1a49aa6-21db-4273-bcfa-bc59b4b9d0dc",
   "metadata": {},
   "source": [
    "### Step 6: Connect to Ray cluster"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "465829b7-f72c-4532-80f5-5a69f470fa18",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2022-03-16 16:30:27,154\tINFO services.py:1412 -- View the Ray dashboard at \u001b[1m\u001b[32mhttp://127.0.0.1:8268\u001b[39m\u001b[22m\n"
     ]
    }
   ],
   "source": [
    "CONNECT_TO_ANYSCALE=False\n",
    "if ray.is_initialized:\n",
    "    ray.shutdown()\n",
    "    if CONNECT_TO_ANYSCALE:\n",
    "        ray.init(address)\n",
    "    else:\n",
    "        ray.init(ignore_reinit_error=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c93caaba-d57b-4a4a-8a06-49dd01474aaf",
   "metadata": {},
   "source": [
    "### Step 7: Run the main Trainer driver"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "d41e4c6c-332b-4c68-8de4-e45c60647e1a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2022-03-16 16:19:11,394\tINFO trainer.py:199 -- Trainer logs will be logged in: /Users/jules/ray_results/train_2022-03-16_16-19-11\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61696)\u001b[0m 2022-03-16 16:19:13,471\tINFO torch.py:66 -- Setting up process group for: env:// [rank=0, world_size=8]\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61697)\u001b[0m 2022-03-16 16:19:13,458\tINFO torch.py:66 -- Setting up process group for: env:// [rank=1, world_size=8]\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61699)\u001b[0m 2022-03-16 16:19:13,466\tINFO torch.py:66 -- Setting up process group for: env:// [rank=2, world_size=8]\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61690)\u001b[0m 2022-03-16 16:19:13,473\tINFO torch.py:66 -- Setting up process group for: env:// [rank=5, world_size=8]\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61690)\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=61692)\u001b[0m 2022-03-16 16:19:13,481\tINFO torch.py:66 -- Setting up process group for: env:// [rank=4, world_size=8]\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61692)\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=61691)\u001b[0m 2022-03-16 16:19:13,465\tINFO torch.py:66 -- Setting up process group for: env:// [rank=3, world_size=8]\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61694)\u001b[0m 2022-03-16 16:19:13,476\tINFO torch.py:66 -- Setting up process group for: env:// [rank=6, world_size=8]\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61694)\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=61693)\u001b[0m 2022-03-16 16:19:13,469\tINFO torch.py:66 -- Setting up process group for: env:// [rank=7, world_size=8]\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61697)\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",
      "2022-03-16 16:19:14,509\tINFO trainer.py:205 -- Run results will be logged in: /Users/jules/ray_results/train_2022-03-16_16-19-11/run_001\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61696)\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=61699)\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=61691)\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=61693)\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=61696)\u001b[0m 2022-03-16 16:19:16,210\tINFO torch.py:244 -- Moving model to device: cpu\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61696)\u001b[0m 2022-03-16 16:19:16,210\tINFO torch.py:247 -- Wrapping provided model in DDP.\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61697)\u001b[0m 2022-03-16 16:19:16,209\tINFO torch.py:244 -- Moving model to device: cpu\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61697)\u001b[0m 2022-03-16 16:19:16,210\tINFO torch.py:247 -- Wrapping provided model in DDP.\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61699)\u001b[0m 2022-03-16 16:19:16,210\tINFO torch.py:244 -- Moving model to device: cpu\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61699)\u001b[0m 2022-03-16 16:19:16,211\tINFO torch.py:247 -- Wrapping provided model in DDP.\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61690)\u001b[0m 2022-03-16 16:19:16,210\tINFO torch.py:244 -- Moving model to device: cpu\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61690)\u001b[0m 2022-03-16 16:19:16,211\tINFO torch.py:247 -- Wrapping provided model in DDP.\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61692)\u001b[0m 2022-03-16 16:19:16,210\tINFO torch.py:244 -- Moving model to device: cpu\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61692)\u001b[0m 2022-03-16 16:19:16,211\tINFO torch.py:247 -- Wrapping provided model in DDP.\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61691)\u001b[0m 2022-03-16 16:19:16,212\tINFO torch.py:244 -- Moving model to device: cpu\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61691)\u001b[0m 2022-03-16 16:19:16,212\tINFO torch.py:247 -- Wrapping provided model in DDP.\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61694)\u001b[0m 2022-03-16 16:19:16,211\tINFO torch.py:244 -- Moving model to device: cpu\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61694)\u001b[0m 2022-03-16 16:19:16,212\tINFO torch.py:247 -- Wrapping provided model in DDP.\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61693)\u001b[0m 2022-03-16 16:19:16,212\tINFO torch.py:244 -- Moving model to device: cpu\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61693)\u001b[0m 2022-03-16 16:19:16,212\tINFO torch.py:247 -- Wrapping provided model in DDP.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61696)\u001b[0m Epoc: 50, Avg validation loss: 1.20, Avg validation accuracy: 58.51%\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61696)\u001b[0m --------------------------------------------------------------------------------\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61697)\u001b[0m Epoc: 50, Avg validation loss: 1.29, Avg validation accuracy: 55.41%\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61697)\u001b[0m --------------------------------------------------------------------------------\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61699)\u001b[0m Epoc: 50, Avg validation loss: 1.26, Avg validation accuracy: 56.29%\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61699)\u001b[0m --------------------------------------------------------------------------------\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61694)\u001b[0m Epoc: 50, Avg validation loss: 1.23, Avg validation accuracy: 56.61%\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61694)\u001b[0m --------------------------------------------------------------------------------\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61691)\u001b[0m Epoc: 50, Avg validation loss: 1.26, Avg validation accuracy: 54.95%\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61691)\u001b[0m --------------------------------------------------------------------------------\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61692)\u001b[0m Epoc: 50, Avg validation loss: 1.26, Avg validation accuracy: 58.06%\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61692)\u001b[0m --------------------------------------------------------------------------------\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61690)\u001b[0m Epoc: 50, Avg validation loss: 1.25, Avg validation accuracy: 56.62%\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61690)\u001b[0m --------------------------------------------------------------------------------\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61693)\u001b[0m Epoc: 50, Avg validation loss: 1.25, Avg validation accuracy: 56.13%\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61693)\u001b[0m --------------------------------------------------------------------------------\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61696)\u001b[0m Epoc: 100, Avg validation loss: 1.14, Avg validation accuracy: 60.95%\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61696)\u001b[0m --------------------------------------------------------------------------------\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61697)\u001b[0m Epoc: 100, Avg validation loss: 1.22, Avg validation accuracy: 57.42%\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61697)\u001b[0m --------------------------------------------------------------------------------\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61699)\u001b[0m Epoc: 100, Avg validation loss: 1.20, Avg validation accuracy: 58.42%\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61699)\u001b[0m --------------------------------------------------------------------------------\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61694)\u001b[0m Epoc: 100, Avg validation loss: 1.17, Avg validation accuracy: 59.43%\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61694)\u001b[0m --------------------------------------------------------------------------------\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61691)\u001b[0m Epoc: 100, Avg validation loss: 1.20, Avg validation accuracy: 57.44%\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61691)\u001b[0m --------------------------------------------------------------------------------\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61692)\u001b[0m Epoc: 100, Avg validation loss: 1.21, Avg validation accuracy: 60.69%\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61692)\u001b[0m --------------------------------------------------------------------------------\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61690)\u001b[0m Epoc: 100, Avg validation loss: 1.19, Avg validation accuracy: 60.00%\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61690)\u001b[0m --------------------------------------------------------------------------------\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61693)\u001b[0m Epoc: 100, Avg validation loss: 1.19, Avg validation accuracy: 57.76%\n",
      "\u001b[2m\u001b[36m(BaseWorkerMixin pid=61693)\u001b[0m --------------------------------------------------------------------------------\n",
      "CPU times: user 2.76 s, sys: 1.64 s, total: 4.4 s\n",
      "Wall time: 7min 51s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "results = train_fashion_mnist(num_workers=number_of_workers, use_gpu=use_gpu)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9c603e8-6ecb-4e81-8204-7098f48d9627",
   "metadata": {},
   "source": [
    "### Step 8: Observe metrics in Tensorboard \n",
    "\n",
    "Subsitute your path `train_path` printed in the cell above"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "babf734e-4e4d-48a4-b20f-b06871effeb6",
   "metadata": {},
   "outputs": [],
   "source": [
    "!tensorboard --logdir ~/ray_results/<train_path>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "f3b66eae-1404-4a36-9bae-0c96c5567627",
   "metadata": {},
   "outputs": [],
   "source": [
    "ray.shutdown()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e89653c-7e00-48ea-b6b1-6e50d9ab4a38",
   "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 `batch_size` and `epochs`\n",
    " 4. Try the second model by chaninge the `mode_type` in `config` to 1\n",
    " 5. Did it improve the accuracy or minimize the loss?\n",
    " 6. Can you try some deep learning regularization techniques to bring the loss down?\n",
    " 7. Change a the loss function and test if that help\n"
   ]
  }
 ],
 "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
}