Train.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training Script for Fast-Pedestrian-Tracking\n",
    "\n",
    "This Python Jupyter notebook re-runs the training process for our project. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Import Modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import  absolute_import\n",
    "# though cupy is not used but without this line, it raise errors...\n",
    "import cupy as cp\n",
    "import os\n",
    "import ipdb\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "from tqdm import tqdm\n",
    "\n",
    "from utils.config import opt\n",
    "from data.dataset import Dataset, TestDataset, inverse_normalize\n",
    "from model.faster_rcnn_vgg16 import FasterRCNNVGG16\n",
    "from torch.utils import data as data_\n",
    "import torch\n",
    "from trainer import FasterRCNNTrainer\n",
    "from utils import array_tool as at\n",
    "from utils.vis_tool import visdom_bbox, vis_bbox\n",
    "from utils.eval_tool import eval_detection_voc\n",
    "\n",
    "# fix for ulimit\n",
    "# https://github.com/pytorch/pytorch/issues/973#issuecomment-346405667\n",
    "import resource\n",
    "\n",
    "rlimit = resource.getrlimit(resource.RLIMIT_NOFILE)\n",
    "resource.setrlimit(resource.RLIMIT_NOFILE, (2048, rlimit[1]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Configurations used"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(opt._parse_all())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Evaluation function\n",
    "Eval is called during training. The function is used to evaluate performance at different points of training. The result is a dictionary consisting of the values for the map (mean average precision) and ap (average precision)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def eval(dataloader, faster_rcnn, test_num=10000):\n",
    "    print(\"\\nEVAL\")\n",
    "    pred_bboxes, pred_labels, pred_scores = list(), list(), list()\n",
    "    gt_bboxes, gt_labels = list(), list()\n",
    "    for ii, (imgs, sizes, gt_bboxes_, gt_labels_) in tqdm(enumerate(dataloader)):\n",
    "        sizes = [sizes[0][0].item(), sizes[1][0].item()]\n",
    "        pred_bboxes_, pred_labels_, pred_scores_ = faster_rcnn.predict(imgs, [sizes])\n",
    "        gt_bboxes += list(gt_bboxes_.numpy())\n",
    "        gt_labels += list(gt_labels_.numpy())\n",
    "        pred_bboxes += pred_bboxes_\n",
    "        pred_labels += pred_labels_\n",
    "        pred_scores += pred_scores_\n",
    "        if ii == test_num: break\n",
    "\n",
    "    result = eval_detection_voc(\n",
    "        pred_bboxes, pred_labels, pred_scores,\n",
    "        gt_bboxes, gt_labels, use_07_metric=True)\n",
    "    return result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define dataloaders and trainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = Dataset(opt)\n",
    "dataloader = data_.DataLoader(dataset, \\\n",
    "                            batch_size=1, \\\n",
    "                            shuffle=True, \\\n",
    "                            # pin_memory=True,\n",
    "                            num_workers=opt.num_workers)\n",
    "\n",
    "valset = TestDataset(opt, split='val')\n",
    "val_dataloader = data_.DataLoader(valset,\n",
    "                                batch_size=1,\n",
    "                                num_workers=opt.test_num_workers,\n",
    "                                shuffle=False, \\\n",
    "                                pin_memory=True\n",
    "                                )\n",
    "\n",
    "print(f\"TRAIN SET: {len(dataloader)} | VAL SET: {len(val_dataloader)} \")\n",
    "print(\"Using Mask VGG\") if opt.mask else print(\"Using normal VGG16\")\n",
    "faster_rcnn = FasterRCNNVGG16(mask=opt.mask)\n",
    "print('model construct completed')\n",
    "trainer = FasterRCNNTrainer(faster_rcnn).cuda()\n",
    "trainer.vis.text(dataset.db.label_names, win='labels')\n",
    "best_map = 0\n",
    "lr_ = opt.lr\n",
    "start_epoch = 0\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training\n",
    "Training runs for a maximum of 13 epochs and the model is evaluated and saved after every epoch on the validation set.\n",
    "The different losses, ground truth images, and predicted images are printed every 100 iterations.\n",
    "\n",
    "The models after each epoch will be saved under 'checkpoints/train_rerun_epoch{epoch_number}'."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.train()\n",
    "for epoch in range(opt.epoch):\n",
    "    trainer.reset_meters()\n",
    "    pbar = tqdm(enumerate(dataloader), total=len(dataloader))\n",
    "    for ii, (img, bbox_, label_, scale) in pbar:\n",
    "        scale = at.scalar(scale)\n",
    "\n",
    "        img, bbox, label = img.cuda().float(), bbox_.cuda(), label_.cuda()\n",
    "        losses = trainer.train_step(img, bbox, label, scale)\n",
    "        if ii % 100 == 0:\n",
    "            rpnloc = losses[0].cpu().data.numpy()\n",
    "            rpncls = losses[1].cpu().data.numpy()\n",
    "            roiloc = losses[2].cpu().data.numpy()\n",
    "            roicls = losses[3].cpu().data.numpy()\n",
    "            tot = losses[4].cpu().data.numpy()\n",
    "            pbar.set_description(f\"Epoch: {epoch} | Batch: {ii} | RPNLoc Loss: {rpnloc:.4f} | RPNclc Loss: {rpncls:.4f} | ROIloc Loss: {roiloc:.4f} | ROIclc Loss: {roicls:.4f} | Total Loss: {tot:.4f}\")\n",
    "        if (ii + 1) % 100 == 0:\n",
    "            print(trainer.get_meter_data())\n",
    "            ori_img_ = inverse_normalize(at.tonumpy(img[0]))\n",
    "            gt_img = vis_bbox(ori_img_,\n",
    "                                 at.tonumpy(bbox_[0]),\n",
    "                                 at.tonumpy(label_[0]))\n",
    "            plt.show()\n",
    "\n",
    "            # plot predicti bboxes\n",
    "            _bboxes, _labels, _scores = trainer.faster_rcnn.predict([ori_img_], visualize=True)\n",
    "            pred_img = vis_bbox(ori_img_,\n",
    "                                   at.tonumpy(_bboxes[0]),\n",
    "                                   at.tonumpy(_labels[0]).reshape(-1),\n",
    "                                   at.tonumpy(_scores[0]))\n",
    "            plt.show()\n",
    "\n",
    "    eval_result = eval(val_dataloader, faster_rcnn, test_num=1000)\n",
    "    lr_ = trainer.faster_rcnn.optimizer.param_groups[0]['lr']\n",
    "    log_info = 'lr:{}, map:{},loss:{}'.format(str(lr_), str(eval_result['map']),\n",
    "                                                      str(trainer.get_meter_data()))\n",
    "    epoch_path = trainer.save(epoch, \n",
    "                              best_map=eval_result['map'], \n",
    "                              model_name=opt.model_name, \n",
    "                              save_path='checkpoints/train_rerun_epoch'+str(epoch))\n",
    "\n",
    "    if eval_result['map'] > best_map:\n",
    "        best_map = eval_result['map']\n",
    "        best_path = epoch_path\n",
    "    if epoch == 9:\n",
    "        trainer.load(best_path)\n",
    "        trainer.faster_rcnn.scale_lr(opt.lr_decay)\n",
    "        lr_ = lr_ * opt.lr_decay\n",
    "\n",
    "    if epoch == 13: \n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}