ssd_train.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import datetime\n",
    "import torch\n",
    "from torch.optim.lr_scheduler import MultiStepLR\n",
    "from torch.utils.data import DataLoader\n",
    "import numpy as np\n",
    "from tqdm.autonotebook import tqdm\n",
    "from pycocotools.cocoeval import COCOeval\n",
    "\n",
    "# custom functions\n",
    "from model import SSD\n",
    "from dataloader import ObjectDataset, collate_fn\n",
    "from utils import Encoder, make_dirs, BoxesGen\n",
    "from utils import SSDAugmentation, Loss\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "device(type='cuda')"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from torch.utils.tensorboard import SummaryWriter\n",
    "# check cuda availability\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "classes = [\"background\", \"car\", \"motorbike\", \"bus\"] \n",
    "colors = [None, (164, 80, 133), (39, 129, 113), (83, 122, 114), (99, 81, 172)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_epoc(modellib, train_loader, epochs, writer, loss_fn, optimizer, scheduler):\n",
    "    train_loss_l = []\n",
    "    train_acc_l = []\n",
    "    for epoch in range(first_epoch, epochs):\n",
    "        modellib.train()\n",
    "        total = 0\n",
    "        sum_loss = 0\n",
    "        corrects = 0\n",
    "        num_iter_per_epoch = len(train_loader)\n",
    "        progress_bar = tqdm(train_loader)\n",
    "        for i, (img, _, _, grd_box, grd_lbl) in enumerate(progress_bar):\n",
    "            cnt = grd_lbl.shape[0]\n",
    "            img = img.cuda()\n",
    "            grd_box = grd_box.cuda()\n",
    "            grd_lbl = grd_lbl.cuda()\n",
    "            modellib.eval()\n",
    "            pred_box, pred_lbl = modellib(img)\n",
    "            pred_box, pred_lbl = pred_box.float(), pred_lbl.float()\n",
    "            # print(\"grd_box:\", grd_box.shape)\n",
    "            # print(\"pred_box:\", pred_box.shape)\n",
    "            # print(\"grd_lbl:\", grd_lbl.shape)\n",
    "            # print(\"pred_lbl:\", pred_lbl.shape)\n",
    "            grd_box = grd_box.transpose(1, 2).contiguous()\n",
    "            loss = loss_fn(pred_box, pred_lbl, grd_box, grd_lbl)\n",
    "\n",
    "            progress_bar.set_description(\"Epoch: {}. Loss: {:.5f}\".format(epoch + 1, loss.item()))\n",
    "            writer.add_scalar(\"Train Loss\", loss.item(), epoch * num_iter_per_epoch + i)\n",
    "\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            scheduler.step()\n",
    "            optimizer.zero_grad()\n",
    "            \n",
    "            \n",
    "            total += cnt\n",
    "            sum_loss += loss.item()\n",
    "            _, pred = torch.max(pred_lbl, 1)\n",
    "            corrects += torch.sum(pred.eq(grd_lbl)).item()\n",
    "        train_loss = sum_loss/total\n",
    "        train_acc = corrects /total\n",
    "        evaluate(model, test_loader, epoch, writer, encoder, nms_threshold)\n",
    "        \n",
    "        # history\n",
    "        train_loss_l.append(train_loss)\n",
    "        train_acc_l.append(train_acc)\n",
    "        print(\"Epoch: {}, Train Loss: {}\".format(epoch, round(train_loss, 3)))\n",
    "        \n",
    "    return modellib, train_loss_l, train_acc_l"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(modellib, test_loader, epoch, writer, encoder, nms_threshold):\n",
    "    modellib.eval()\n",
    "    detections = []\n",
    "    cat_ids = test_loader.dataset.coco.getCatIds()\n",
    "    progress_bar = tqdm(test_loader)\n",
    "    count = 0\n",
    "    for nbatch, (img, img_id, img_size, grd_box, grd_lbl) in enumerate(progress_bar):\n",
    "        progress_bar.set_description(\"Validation : {}\".format(nbatch))\n",
    "        img = img.cuda()\n",
    "        with torch.no_grad():\n",
    "            # Get predictions\n",
    "            pred_box, pred_lbl = modellib(img)\n",
    "            pred_box, pred_lbl = pred_box.float(), pred_lbl.float()\n",
    "            \n",
    "            # print(\"grd_box:\", grd_box.shape)\n",
    "            # print(\"pred_box:\", pred_box.shape)\n",
    "            # print(\"grd_lbl:\", grd_lbl.shape)\n",
    "            # print(\"pred_lbl:\", pred_lbl.shape)       \n",
    "        \n",
    "            for idx in range(pred_box.shape[0]):\n",
    "                pred_box_i = pred_box[idx, :, :].unsqueeze(0)\n",
    "                pred_lbl_i = pred_lbl[idx, :, :].unsqueeze(0)\n",
    "                try:\n",
    "                    result = encoder.decode_batch(pred_box_i, pred_lbl_i, nms_threshold, 200)[0]\n",
    "                except:\n",
    "                    print(\"No object detected in idx: {}\".format(idx))\n",
    "                    continue\n",
    "\n",
    "                height, width = img_size[idx]\n",
    "                loc, label, prob = [r.cpu().numpy() for r in result]\n",
    "                for loc_, label_, prob_ in zip(loc, label, prob):\n",
    "                    detections.append([img_id[idx], loc_[0] * width, loc_[1] * height, (loc_[2] - loc_[0]) * width,\n",
    "                                       (loc_[3] - loc_[1]) * height, prob_,\n",
    "                                       cat_ids[label_ - 1]])\n",
    "        count = count + 1\n",
    "\n",
    "    detections = np.array(detections, dtype=np.float32)\n",
    "    coco_eval = COCOeval(test_loader.dataset.coco, test_loader.dataset.coco.loadRes(detections), iouType=\"bbox\")\n",
    "    coco_eval.evaluate()\n",
    "    coco_eval.accumulate()\n",
    "    coco_eval.summarize()\n",
    "    checkpoint = {\"epoch\": epoch, \"model_state_dict\": model.state_dict(),\n",
    "                \"optimizer\": optimizer.state_dict(), \n",
    "                \"scheduler\": scheduler.state_dict()}\n",
    "    torch.save(checkpoint, checkpoint_path)\n",
    "\n",
    "    writer.add_scalar(\"Test mAP\", coco_eval.stats[0], epoch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.cuda.manual_seed(123)\n",
    "batch_size = 1\n",
    "num_workers = 4\n",
    "log_dir = \"logs\"\n",
    "lr = 2.6e-3\n",
    "epochs = 10\n",
    "momentum = 0.9\n",
    "nms_threshold = 0.5\n",
    "weight_decay = 0.0005\n",
    "multi_step = [43, 54]\n",
    "data_path = \"data/vehicles\"\n",
    "save_folder = \"trained_models\"\n",
    "date_time = datetime.datetime.now().strftime(\"%d-%m-%Y_%H-%M-%S\")\n",
    "   \n",
    "log_path = os.path.join(log_dir, date_time)\n",
    "save_path = os.path.join(log_path, save_folder)\n",
    "make_dirs(save_path)\n",
    "checkpoint_path = os.path.join(save_path, \"model_ssd.pth\")\n",
    "\n",
    "train_params = {\"batch_size\": batch_size,  \"shuffle\": True,\n",
    "                \"drop_last\": False, \"num_workers\": num_workers,\n",
    "                \"collate_fn\": collate_fn}\n",
    "\n",
    "test_params = {\"batch_size\": batch_size, \"shuffle\": False,\n",
    "                \"drop_last\": False, \"num_workers\": num_workers,\n",
    "                \"collate_fn\": collate_fn}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading annotations into memory...\n",
      "Done (t=0.16s)\n",
      "creating index...\n",
      "index created!\n",
      "loading annotations into memory...\n",
      "Done (t=0.03s)\n",
      "creating index...\n",
      "index created!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/jyothir-dr/miniconda3/envs/detect/lib/python3.8/site-packages/torch/nn/_reduction.py:42: UserWarning: size_average and reduce args will be deprecated, please use reduction='none' instead.\n",
      "  warnings.warn(warning.format(ret))\n"
     ]
    }
   ],
   "source": [
    "dboxes = BoxesGen()\n",
    "encoder = Encoder(dboxes)\n",
    "\n",
    "model = SSD(num_classes=4).to(device)\n",
    "    \n",
    "train_set = ObjectDataset(data_path, \"train\", SSDAugmentation(dboxes, (300, 300), val=False))\n",
    "train_loader = DataLoader(train_set, **train_params)\n",
    "test_set = ObjectDataset(data_path, \"val\", SSDAugmentation(dboxes, (300, 300), val=True))\n",
    "test_loader = DataLoader(test_set, **test_params)\n",
    "\n",
    "# update lr\n",
    "lr = lr * (batch_size / 32)\n",
    "\n",
    "loss_fn = Loss(dboxes).to(device)\n",
    "\n",
    "optimizer = torch.optim.SGD(model.parameters(), lr=lr, momentum=momentum,\n",
    "                            weight_decay=weight_decay, nesterov=True)\n",
    "\n",
    "scheduler = MultiStepLR(optimizer = optimizer, milestones = multi_step, gamma=0.1)\n",
    "# model.cuda()\n",
    "# loss_fn.cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "if os.path.isfile(checkpoint_path):\n",
    "    checkpoint = torch.load(checkpoint_path)\n",
    "    first_epoch = checkpoint[\"epoch\"] + 1\n",
    "    # model.module.load_state_dict(checkpoint[\"model_state_dict\"])\n",
    "    model.load_state_dict(checkpoint[\"model_state_dict\"])\n",
    "\n",
    "    scheduler.load_state_dict(checkpoint[\"scheduler\"])\n",
    "    optimizer.load_state_dict(checkpoint[\"optimizer\"])\n",
    "else:\n",
    "    first_epoch = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2022-07-28 19:02:37.602746: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library libcudart.so.11.0\n"
     ]
    }
   ],
   "source": [
    "writer = SummaryWriter(log_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fad95b29486a49c3b9af7c64acc191f6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1524 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2fe69cc684a845c9842cb29745b9cdba",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/344 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading and preparing results...\n",
      "Converting ndarray to lists...\n",
      "(68800, 7)\n",
      "0/68800\n",
      "DONE (t=0.60s)\n",
      "creating index...\n",
      "index created!\n",
      "Running per image evaluation...\n",
      "Evaluate annotation type *bbox*\n",
      "DONE (t=4.29s).\n",
      "Accumulating evaluation results...\n",
      "DONE (t=0.63s).\n",
      " Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.001\n",
      " Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 0.001\n",
      " Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.001\n",
      " Average Precision  (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.000\n",
      " Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.001\n",
      " Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.000\n",
      " Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.000\n",
      " Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.002\n",
      " Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.007\n",
      " Average Recall     (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.000\n",
      " Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.009\n",
      " Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.007\n",
      "Epoch: 0, Train Loss: 9.774\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5da7530d66b04b98a65664b657911e41",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1524 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c35c40e36ad947a084f050c075c1fe7a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/344 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading and preparing results...\n",
      "Converting ndarray to lists...\n",
      "(68800, 7)\n",
      "0/68800\n",
      "DONE (t=0.49s)\n",
      "creating index...\n",
      "index created!\n",
      "Running per image evaluation...\n",
      "Evaluate annotation type *bbox*\n",
      "DONE (t=4.26s).\n",
      "Accumulating evaluation results...\n",
      "DONE (t=0.71s).\n",
      " Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.001\n",
      " Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 0.001\n",
      " Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.001\n",
      " Average Precision  (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.000\n",
      " Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.001\n",
      " Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.000\n",
      " Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.001\n",
      " Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.002\n",
      " Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.009\n",
      " Average Recall     (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.000\n",
      " Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.012\n",
      " Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.009\n",
      "Epoch: 1, Train Loss: 9.699\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "35ee9241e9764e69a03beb27414ea45e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1524 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4c5691c06beb4110b37d1b93e716d05f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/344 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading and preparing results...\n",
      "Converting ndarray to lists...\n",
      "(68800, 7)\n",
      "0/68800\n",
      "DONE (t=0.47s)\n",
      "creating index...\n",
      "index created!\n",
      "Running per image evaluation...\n",
      "Evaluate annotation type *bbox*\n",
      "DONE (t=4.73s).\n",
      "Accumulating evaluation results...\n",
      "DONE (t=0.70s).\n",
      " Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.000\n",
      " Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 0.001\n",
      " Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.000\n",
      " Average Precision  (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.000\n",
      " Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.001\n",
      " Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.000\n",
      " Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.001\n",
      " Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.003\n",
      " Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.010\n",
      " Average Recall     (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.000\n",
      " Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.014\n",
      " Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.011\n",
      "Epoch: 2, Train Loss: 9.662\n"
     ]
    }
   ],
   "source": [
    "epochs = 3\n",
    "modellib, train_loss_l, train_acc_l = train_epoc(model, train_loader, epochs, writer, loss_fn, optimizer, scheduler)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.13 ('detect')",
   "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.13"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "2cbf2d2672eb1d51db66da973b7f534e15d3112a510a2d9edb972ac6185763d1"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}