lecture_4/detection.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2739b930",
   "metadata": {},
   "source": [
    "# Let's build an AI to detect robots!\n",
    "\n",
    "Ceylon & HelloElwin @ HKU Astar"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f32be3e2",
   "metadata": {},
   "source": [
    "## Installation\n",
    "\n",
    "### MacOS\n",
    "```shell\n",
    "pip3 install torch torchvision\n",
    "pip3 install notebook\n",
    "```\n",
    "\n",
    "### Windows\n",
    "```shell\n",
    "pip3 install torch==1.8.1+cpu torchvision==0.9.1+cpu -f https://download.pytorch.org/whl/torch_stable.html\n",
    "pip3 install notebook\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd7d336f",
   "metadata": {},
   "source": [
    "## Basic workflows\n",
    "* working with data; \n",
    "* creating models;\n",
    "* optimizing model parameters;\n",
    "* testing the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "896852f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "from torchvision.io import read_image\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "import torch.nn.functional as F\n",
    "import os\n",
    "\n",
    "# Size of images\n",
    "IMG_R = 48\n",
    "IMG_C = 64\n",
    "\n",
    "BATCH_SIZE = 8\n",
    "EPOCH = 20"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bfb65383",
   "metadata": {},
   "source": [
    "## Working with data\n",
    "\n",
    "PyTorch has two primitives to work with data: `torch.utils.data.DataLoader` and `torch.utils.data.Dataset`. \n",
    "`Dataset` stores the samples and their corresponding labels, and `DataLoader` wraps an iterable around a `Dataset`.\n",
    "\n",
    "Required member function of a `Dataset` class:\n",
    "\n",
    "- `__init__`\n",
    "\n",
    "    The `__init__` function is run once when instantiating the `RobotDataset` object. We initialize the directory with putting all images and their corresponding labels into lists.\n",
    "\n",
    "- `__len__`\n",
    "    The `__len__` function returns the number of samples in our dataset.\n",
    "\n",
    "- `__getitem__`\n",
    "\n",
    "    The `__getitem__` function loads and returns a sample from the dataset at the given index `idx`. \n",
    "\n",
    "Some external functions we use:\n",
    "\n",
    "- `os.listdir`: get the list of names of all files and directories in the specified directory.\n",
    "\n",
    "- `read_image`: converts an image to a `Tensor`, a datatype for multi-dimensional matrices in PyTorch.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "79111e7a",
   "metadata": {},
   "outputs": [],
   "source": [
    "class RobotDataset(Dataset):\n",
    "    def __init__(self, robots_dir, others_dir):\n",
    "        self.images = []\n",
    "        self.labels = []\n",
    "        for file_name in os.listdir(robots_dir):\n",
    "            image = read_image(robots_dir + file_name)\n",
    "            image = image[:, ::10, ::10]\n",
    "            self.images.append(image)\n",
    "            self.labels.append(0)\n",
    "        for file_name in os.listdir(others_dir):\n",
    "            image = read_image(others_dir + file_name)\n",
    "            self.images.append(image)\n",
    "            self.labels.append(1)\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.images)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return self.images[idx], self.labels[idx]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41c98530",
   "metadata": {},
   "source": [
    "To accelerate operations like matrix multiplications in the neural network, we allocate the whole model and data to the GPU if available.\n",
    "\n",
    "This is done by member function `torch.Tensor.to()` of a `Tensor`. We specify a `device`, e.g. `cuda` (GPU) or `cpu`, and then use `x.to(device)` to move the Tensor `x` to that `device`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61005c20",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data = RobotDataset('./datasets/train/robots/', './datasets/train/others/')\n",
    "test_data = RobotDataset('./datasets/test/robots/', './datasets/test/others/')\n",
    "train_dataloader = DataLoader(train_data, batch_size=BATCH_SIZE, shuffle=True)\n",
    "test_dataloader = DataLoader(test_data, batch_size=BATCH_SIZE, shuffle=False)\n",
    "\n",
    "# Here we select a device for later use\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e94e155b",
   "metadata": {},
   "source": [
    "## Creating models\n",
    "* To define a neural network or any ML models in PyTorch, we create a class that inherits from `nn.Module`. \n",
    "* We define the layers of the network in the `__init__` function and specify how data will pass through the network in the `forward` function. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "bde5ce65",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MLP(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__() \n",
    "        self.n = nn.Sequential(\n",
    "            nn.Linear(3 * IMG_R * IMG_C, 4608),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(4608, 1152),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(1152, 288),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(288, 64),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(64, 2)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        flatten = nn.Flatten()\n",
    "        x = flatten(x)\n",
    "        x = self.n(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "043e4c8c",
   "metadata": {},
   "outputs": [],
   "source": [
    "class CNN(nn.Module):\n",
    "\n",
    "    def __init__(self):\n",
    "        super(CNN, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(3, 6, 4)\n",
    "        self.conv2 = nn.Conv2d(6, 16, 4)\n",
    "        self.fc1 = nn.Linear(1872, 256)\n",
    "        self.fc2 = nn.Linear(256, 128)\n",
    "        self.fc3 = nn.Linear(128, 2)\n",
    "\n",
    "    def forward(self, x):\n",
    "        # Max pooling over a (2, 2) window\n",
    "        x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))\n",
    "        # If the size is a square, you can specify with a single number\n",
    "        x = F.max_pool2d(F.relu(self.conv2(x)), 2)\n",
    "        \n",
    "        x = torch.flatten(x, 1) # flatten all dimensions except the batch dimension\n",
    "        x = F.relu(self.fc1(x))\n",
    "        x = F.relu(self.fc2(x))\n",
    "        x = self.fc3(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0ee2cb7",
   "metadata": {},
   "source": [
    "## Optimizing model parameters (train)\n",
    "\n",
    "We iteratively train the model through all data in the `dataloader` in each `epoch`, and use a certain kind of `optimizer` to do the \"Gradient Descent\" for us."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "c46ac68f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(dataloader, model, loss_fn, optimizer):\n",
    "    model.train()\n",
    "    size = len(dataloader.dataset)\n",
    "    for epoch in range(EPOCH):\n",
    "        epoch_loss = 0\n",
    "        for batch, (x, y) in enumerate(dataloader):\n",
    "            x, y = x.to(device), y.to(device)\n",
    "\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "            pred = model(x.float())\n",
    "            loss = loss_fn(pred, y)\n",
    "\n",
    "            # Backpropagation\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "            loss, current = loss.item(), min(batch * BATCH_SIZE, size)\n",
    "            print(f\"loss: {loss:.6f} [{current}/{size}]   \", end=\"\\r\")\n",
    "            \n",
    "            epoch_loss += loss\n",
    "            \n",
    "        print(f\"Epoch [{epoch}/{EPOCH}]: Loss={epoch_loss:.6f}        \")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "99bf4b7b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [0/20]: Loss=38.207765        \n",
      "Epoch [1/20]: Loss=19.360587        \n",
      "Epoch [2/20]: Loss=14.358377        \n",
      "Epoch [3/20]: Loss=10.978099        \n",
      "Epoch [4/20]: Loss=8.510617        \n",
      "Epoch [5/20]: Loss=6.641623        \n",
      "Epoch [6/20]: Loss=5.147152        \n",
      "Epoch [7/20]: Loss=3.923710        \n",
      "Epoch [8/20]: Loss=2.923370        \n",
      "Epoch [9/20]: Loss=2.117846        \n",
      "Epoch [10/20]: Loss=1.513425        \n",
      "Epoch [11/20]: Loss=1.086297        \n",
      "Epoch [12/20]: Loss=0.792458        \n",
      "Epoch [13/20]: Loss=0.589680        \n",
      "Epoch [14/20]: Loss=0.452154        \n",
      "Epoch [15/20]: Loss=0.357078        \n",
      "Epoch [16/20]: Loss=0.287551        \n",
      "Epoch [17/20]: Loss=0.237379        \n",
      "Epoch [18/20]: Loss=0.199184        \n",
      "Epoch [19/20]: Loss=0.169702        \n"
     ]
    }
   ],
   "source": [
    "model = CNN()\n",
    "loss_fn = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)\n",
    "model = model.to(device)\n",
    "train(train_dataloader, model, loss_fn, optimizer)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e787a16",
   "metadata": {},
   "source": [
    "## Testing the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a84a23e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def test(dataloader, model, loss_fn):\n",
    "    model.eval()\n",
    "    size = len(dataloader.dataset)\n",
    "    test_loss = 0\n",
    "    \n",
    "    for batch, (x, y) in enumerate(dataloader):\n",
    "        x, y = x.to(device), y.to(device)\n",
    "\n",
    "        pred = model(x.float())\n",
    "        loss = loss_fn(pred, y)\n",
    "\n",
    "        loss, current = loss.item(), min(batch * BATCH_SIZE, size)\n",
    "        print(f\"loss: {loss:>7f} [{current}/{size}]   \", end=\"\\r\")\n",
    "\n",
    "        test_loss += loss\n",
    "\n",
    "    print(f\"Test: Loss={test_loss:.6f}              \")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "f9bfe94d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss: 0.000021 [0/158]   \r\n",
      "loss: 0.000013 [8/158]   \r\n",
      "loss: 0.000044 [16/158]   \r\n",
      "loss: 0.000009 [24/158]   \r\n",
      "loss: 0.000043 [32/158]   \r\n",
      "loss: 0.000800 [40/158]   \r\n",
      "loss: 0.000005 [48/158]   \r\n",
      "loss: 0.000313 [56/158]   \r\n",
      "loss: 0.002030 [64/158]   \r\n",
      "loss: 0.000953 [72/158]   \r\n",
      "loss: 0.003836 [80/158]   \r\n",
      "loss: 0.003799 [88/158]   \r\n",
      "loss: 0.003838 [96/158]   \r\n",
      "loss: 0.003798 [104/158]   \r\n",
      "loss: 0.003790 [112/158]   \r\n",
      "loss: 0.004637 [120/158]   \r\n",
      "loss: 0.003808 [128/158]   \r\n",
      "loss: 0.003799 [136/158]   \r\n",
      "loss: 0.003833 [144/158]   \r\n",
      "loss: 0.003803 [152/158]   \r\n",
      "Test: Loss=0.043172              \n"
     ]
    }
   ],
   "source": [
    "test(test_dataloader, model, loss_fn)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6bf0a477",
   "metadata": {},
   "source": [
    "## Useful Links & Reference\n",
    "https://pytorch.org/tutorials/"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.12 64-bit",
   "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.9.12"
  },
  "vscode": {
   "interpreter": {
    "hash": "b0fa6594d8f4cbf19f97940f81e996739fb7646882a419484c72d19e05852a7e"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}