cv-competition-1/pytorch_baseline/pytorch_baseline.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "running build_ext\n",
      "skipping 'compute_overlap.c' Cython extension (up-to-date)\n",
      "copying build/lib.linux-x86_64-3.8/compute_overlap.cpython-38-x86_64-linux-gnu.so -> \n"
     ]
    }
   ],
   "source": [
    "!python setup.py build_ext --inplace\n",
    "\n",
    "import os\n",
    "import random\n",
    "import math\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from PIL import Image, ImageDraw\n",
    "\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.optim import lr_scheduler\n",
    "\n",
    "import torchvision\n",
    "from torchvision import datasets, models, transforms\n",
    "from torchvision.models.detection.retinanet import RetinaNet\n",
    "from torchvision.models.detection.faster_rcnn import FasterRCNN\n",
    "\n",
    "import  torchvision.transforms.functional as F\n",
    "\n",
    "from typing import Any, Callable, Dict, List, Optional, Tuple, Union\n",
    "import xml.etree.ElementTree as ET\n",
    "import collections\n",
    "from torchvision.datasets.voc import VisionDataset\n",
    "\n",
    "from metrics import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import defaultdict, deque\n",
    "import time\n",
    "import datetime\n",
    "\n",
    "import torch.distributed as dist\n",
    "import  torchvision.transforms.functional as F\n",
    "\n",
    "class SmoothedValue(object):\n",
    "    \"\"\"Track a series of values and provide access to smoothed values over a\n",
    "    window or the global series average.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, window_size=20, fmt=None):\n",
    "        if fmt is None:\n",
    "            fmt = \"{median:.4f} ({global_avg:.4f})\"\n",
    "        self.deque = deque(maxlen=window_size)\n",
    "        self.total = 0.0\n",
    "        self.count = 0\n",
    "        self.fmt = fmt\n",
    "\n",
    "    def update(self, value, n=1):\n",
    "        self.deque.append(value)\n",
    "        self.count += n\n",
    "        self.total += value * n\n",
    "\n",
    "    def synchronize_between_processes(self):\n",
    "        \"\"\"\n",
    "        Warning: does not synchronize the deque!\n",
    "        \"\"\"\n",
    "        if not is_dist_avail_and_initialized():\n",
    "            return\n",
    "        t = torch.tensor([self.count, self.total], dtype=torch.float64, device='cuda')\n",
    "        dist.barrier()\n",
    "        dist.all_reduce(t)\n",
    "        t = t.tolist()\n",
    "        self.count = int(t[0])\n",
    "        self.total = t[1]\n",
    "\n",
    "    @property\n",
    "    def median(self):\n",
    "        d = torch.tensor(list(self.deque))\n",
    "        return d.median().item()\n",
    "\n",
    "    @property\n",
    "    def avg(self):\n",
    "        d = torch.tensor(list(self.deque), dtype=torch.float32)\n",
    "        return d.mean().item()\n",
    "\n",
    "    @property\n",
    "    def global_avg(self):\n",
    "        return self.total / self.count\n",
    "\n",
    "    @property\n",
    "    def max(self):\n",
    "        return max(self.deque)\n",
    "\n",
    "    @property\n",
    "    def value(self):\n",
    "        return self.deque[-1]\n",
    "\n",
    "    def __str__(self):\n",
    "        return self.fmt.format(\n",
    "            median=self.median,\n",
    "            avg=self.avg,\n",
    "            global_avg=self.global_avg,\n",
    "            max=self.max,\n",
    "            value=self.value)\n",
    "\n",
    "class MetricLogger(object):\n",
    "    def __init__(self, delimiter=\"\\t\"):\n",
    "        self.meters = defaultdict(SmoothedValue)\n",
    "        self.delimiter = delimiter\n",
    "\n",
    "    def update(self, **kwargs):\n",
    "        for k, v in kwargs.items():\n",
    "            if isinstance(v, torch.Tensor):\n",
    "                v = v.item()\n",
    "            assert isinstance(v, (float, int))\n",
    "            self.meters[k].update(v)\n",
    "\n",
    "    def __getattr__(self, attr):\n",
    "        if attr in self.meters:\n",
    "            return self.meters[attr]\n",
    "        if attr in self.__dict__:\n",
    "            return self.__dict__[attr]\n",
    "        raise AttributeError(\"'{}' object has no attribute '{}'\".format(\n",
    "            type(self).__name__, attr))\n",
    "\n",
    "    def __str__(self):\n",
    "        loss_str = []\n",
    "        for name, meter in self.meters.items():\n",
    "            loss_str.append(\n",
    "                \"{}: {}\".format(name, str(meter))\n",
    "            )\n",
    "        return self.delimiter.join(loss_str)\n",
    "\n",
    "    def synchronize_between_processes(self):\n",
    "        for meter in self.meters.values():\n",
    "            meter.synchronize_between_processes()\n",
    "\n",
    "    def add_meter(self, name, meter):\n",
    "        self.meters[name] = meter\n",
    "\n",
    "    def log_every(self, iterable, print_freq, header=None):\n",
    "        i = 0\n",
    "        if not header:\n",
    "            header = ''\n",
    "        start_time = time.time()\n",
    "        end = time.time()\n",
    "        iter_time = SmoothedValue(fmt='{avg:.4f}')\n",
    "        data_time = SmoothedValue(fmt='{avg:.4f}')\n",
    "        space_fmt = ':' + str(len(str(len(iterable)))) + 'd'\n",
    "        if torch.cuda.is_available():\n",
    "            log_msg = self.delimiter.join([\n",
    "                header,\n",
    "                '[{0' + space_fmt + '}/{1}]',\n",
    "                'eta: {eta}',\n",
    "                '{meters}',\n",
    "                'time: {time}',\n",
    "                'data: {data}',\n",
    "                'max mem: {memory:.0f}'\n",
    "            ])\n",
    "        else:\n",
    "            log_msg = self.delimiter.join([\n",
    "                header,\n",
    "                '[{0' + space_fmt + '}/{1}]',\n",
    "                'eta: {eta}',\n",
    "                '{meters}',\n",
    "                'time: {time}',\n",
    "                'data: {data}'\n",
    "            ])\n",
    "        MB = 1024.0 * 1024.0\n",
    "        for obj in iterable:\n",
    "            data_time.update(time.time() - end)\n",
    "            yield obj\n",
    "            iter_time.update(time.time() - end)\n",
    "            if i % print_freq == 0 or i == len(iterable) - 1:\n",
    "                eta_seconds = iter_time.global_avg * (len(iterable) - i)\n",
    "                eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))\n",
    "                if torch.cuda.is_available():\n",
    "                    print(log_msg.format(\n",
    "                        i, len(iterable), eta=eta_string,\n",
    "                        meters=str(self),\n",
    "                        time=str(iter_time), data=str(data_time),\n",
    "                        memory=torch.cuda.max_memory_allocated() / MB))\n",
    "                else:\n",
    "                    print(log_msg.format(\n",
    "                        i, len(iterable), eta=eta_string,\n",
    "                        meters=str(self),\n",
    "                        time=str(iter_time), data=str(data_time)))\n",
    "            i += 1\n",
    "            end = time.time()\n",
    "        total_time = time.time() - start_time\n",
    "        total_time_str = str(datetime.timedelta(seconds=int(total_time)))\n",
    "        print('{} Total time: {} ({:.4f} s / it)'.format(\n",
    "            header, total_time_str, total_time / len(iterable)))\n",
    "        \n",
    "        \n",
    "\n",
    "def warmup_lr_scheduler(optimizer, warmup_iters, warmup_factor):\n",
    "\n",
    "    def f(x):\n",
    "        if x >= warmup_iters:\n",
    "            return 1\n",
    "        alpha = float(x) / warmup_iters\n",
    "        return warmup_factor * (1 - alpha) + alpha\n",
    "\n",
    "    return torch.optim.lr_scheduler.LambdaLR(optimizer, f)\n",
    "\n",
    "def is_dist_avail_and_initialized():\n",
    "    if not dist.is_available():\n",
    "        return False\n",
    "    if not dist.is_initialized():\n",
    "        return False\n",
    "    return True\n",
    "\n",
    "\n",
    "def get_world_size():\n",
    "    if not is_dist_avail_and_initialized():\n",
    "        return 1\n",
    "    return dist.get_world_size()\n",
    "\n",
    "def reduce_dict(input_dict, average=True):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        input_dict (dict): all the values will be reduced\n",
    "        average (bool): whether to do average or sum\n",
    "    Reduce the values in the dictionary from all processes so that all processes\n",
    "    have the averaged results. Returns a dict with the same fields as\n",
    "    input_dict, after reduction.\n",
    "    \"\"\"\n",
    "    world_size = get_world_size()\n",
    "    if world_size < 2:\n",
    "        return input_dict\n",
    "    with torch.no_grad():\n",
    "        names = []\n",
    "        values = []\n",
    "        # sort the keys so that they are consistent across processes\n",
    "        for k in sorted(input_dict.keys()):\n",
    "            names.append(k)\n",
    "            values.append(input_dict[k])\n",
    "        values = torch.stack(values, dim=0)\n",
    "        dist.all_reduce(values)\n",
    "        if average:\n",
    "            values /= world_size\n",
    "        reduced_dict = {k: v for k, v in zip(names, values)}\n",
    "    return reduced_dict\n",
    "\n",
    "def train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq):\n",
    "    model.train()\n",
    "    metric_logger = MetricLogger(delimiter=\"  \")\n",
    "    metric_logger.add_meter('lr', SmoothedValue(window_size=1, fmt='{value:.6f}'))\n",
    "    header = 'Epoch: [{}]'.format(epoch)\n",
    "\n",
    "    lr_scheduler = None\n",
    "    if epoch == 0:\n",
    "        warmup_factor = 1. / 1000\n",
    "        warmup_iters = min(1000, len(data_loader) - 1)\n",
    "\n",
    "        lr_scheduler = warmup_lr_scheduler(optimizer, warmup_iters, warmup_factor)\n",
    "\n",
    "    for images, targets in metric_logger.log_every(data_loader, print_freq, header):\n",
    "        images = list(image.to(device) for image in images)\n",
    "        targets = [{k: v.to(device) for k, v in t.items()} for t in targets]\n",
    "\n",
    "        loss_dict = model(images, targets)\n",
    "\n",
    "        losses = sum(loss for loss in loss_dict.values())\n",
    "\n",
    "        # reduce losses over all GPUs for logging purposes\n",
    "        loss_dict_reduced = reduce_dict(loss_dict)\n",
    "        losses_reduced = sum(loss for loss in loss_dict_reduced.values())\n",
    "\n",
    "        loss_value = losses_reduced.item()\n",
    "\n",
    "        if not math.isfinite(loss_value):\n",
    "            print(\"Loss is {}, stopping training\".format(loss_value))\n",
    "            print(loss_dict_reduced)\n",
    "            sys.exit(1)\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        losses.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        if lr_scheduler is not None:\n",
    "            lr_scheduler.step()\n",
    "\n",
    "        metric_logger.update(loss=losses_reduced, **loss_dict_reduced)\n",
    "        metric_logger.update(lr=optimizer.param_groups[0][\"lr\"])\n",
    "\n",
    "    return metric_logger\n",
    "\n",
    "@torch.no_grad()\n",
    "def evaluate(model, data_loader, device='cuda'):\n",
    "    n_threads = torch.get_num_threads()\n",
    "    torch.set_num_threads(1) # Does it nessesary? Who knows...\n",
    "    cpu_device = torch.device(\"cpu\")\n",
    "    inference_res = []\n",
    "    model.eval()\n",
    "    \n",
    "    for images, targets in data_loader:\n",
    "        images = list(img.to(device) for img in images)\n",
    "\n",
    "        if torch.cuda.is_available():\n",
    "            torch.cuda.synchronize()\n",
    "        outputs = model(images)\n",
    "\n",
    "        outputs = [{k: v.to(cpu_device) for k, v in t.items()} for t in outputs]\n",
    "        res = targets, outputs\n",
    "        inference_res.append(res)\n",
    "\n",
    "    torch.set_num_threads(n_threads)\n",
    "    return inference_res\n",
    "\n",
    "\n",
    "def collate_fn(batch):\n",
    "    \"\"\"\n",
    "    This function helps when we have different number of object instances\n",
    "    in the batches in the dataset.\n",
    "    \"\"\"\n",
    "    return tuple(zip(*batch))\n",
    "\n",
    "class ToTensor(object):\n",
    "    def __call__(self, image, target):\n",
    "        image = F.to_tensor(image)\n",
    "        return image, target\n",
    "\n",
    "class Compose(object):\n",
    "    def __init__(self, transforms):\n",
    "        self.transforms = transforms\n",
    "\n",
    "    def __call__(self, image, target):\n",
    "        for t in self.transforms:\n",
    "            image, target = t(image, target)\n",
    "        return image, target\n",
    "    \n",
    "class RandomHorizontalFlip(object):\n",
    "    def __init__(self, prob):\n",
    "        self.prob = prob\n",
    "\n",
    "    def __call__(self, image, target):\n",
    "        if random.random() < self.prob:\n",
    "            height, width = image.shape[-2:]\n",
    "            image = image.flip(-1)\n",
    "            bbox = target[\"boxes\"]\n",
    "            if bbox.shape[0]>0:\n",
    "                bbox[:, [0, 2]] = width - bbox[:, [2, 0]]\n",
    "                target[\"boxes\"] = bbox\n",
    "        return image, target\n",
    "    \n",
    "class Resize(object):\n",
    "    def __init__(self, target_size):\n",
    "        self.target_size = target_size\n",
    "    \n",
    "    def __call__(self, image, target):\n",
    "        old_height, old_width = image.shape[-2:]\n",
    "        image = F.resize(image, self.target_size, interpolation=F.InterpolationMode.BILINEAR)\n",
    "        new_height, new_width = image.shape[-2:]\n",
    "        bbox = target[\"boxes\"]\n",
    "        if bbox.shape[0]>0:\n",
    "            bbox[:, [0, 2]] = bbox[:, [0, 2]] * new_width / old_width\n",
    "            bbox[:, [1, 3]] = bbox[:, [1, 3]] * new_height / old_height\n",
    "            target[\"boxes\"] = bbox\n",
    "        \n",
    "        \n",
    "        return image, target \n",
    "    \n",
    "def get_transform(train, target_size):\n",
    "    transforms = []\n",
    "    # converts the image, a PIL image, into a PyTorch Tensor\n",
    "    transforms.append(ToTensor())\n",
    "    transforms.append(Resize(target_size)) \n",
    "    if train:\n",
    "        # during training, randomly flip the training images\n",
    "        # and ground-truth for data augmentation\n",
    "        transforms.append(RandomHorizontalFlip(0.5))\n",
    "    return Compose(transforms)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "params = {}\n",
    "params['target_size']=(1500,2000)\n",
    "params['batch_size'] = 1\n",
    "params['lr'] = 0.001\n",
    "\n",
    "voc_root = '../data/TrainingData'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Reworked class from pytorch (see https://pytorch.org/vision/0.8/_modules/torchvision/datasets/voc.html#VOCDetection)\n",
    "\n",
    "class LADDDataSET(torchvision.datasets.VisionDataset):\n",
    "    def __init__(\n",
    "            self,\n",
    "            root: str,\n",
    "            image_set: str,\n",
    "            transforms: Optional[Callable] = None):     \n",
    "        super(LADDDataSET, self).__init__(root, transforms=transforms)\n",
    "        self.image_set = image_set\n",
    "\n",
    "        voc_root = root\n",
    "        image_dir = os.path.join(voc_root, 'JPEGImages')\n",
    "        annotation_dir = os.path.join(voc_root, 'Annotations')\n",
    "\n",
    "        if not os.path.isdir(voc_root):\n",
    "            raise RuntimeError('Dataset not found or corrupted.')\n",
    "\n",
    "        splits_dir = os.path.join(voc_root, 'ImageSets/Main')\n",
    "        split_f = os.path.join(splits_dir, image_set.rstrip('\\n') + '.txt')\n",
    "\n",
    "        with open(os.path.join(split_f), \"r\") as f:\n",
    "            file_names = [x.strip() for x in f.readlines()]\n",
    "\n",
    "        self.images = [os.path.join(image_dir, x + \".jpg\") for x in file_names]\n",
    "        self.annotations = [os.path.join(annotation_dir, x + \".xml\") for x in file_names]\n",
    "        assert (len(self.images) == len(self.annotations))\n",
    "        \n",
    "    def __getitem__(self, index: int) -> Tuple[Any, Any]:\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            index (int): Index\n",
    "\n",
    "        Returns:\n",
    "            tuple: (image, target) where target is a dictionary of the XML tree.\n",
    "        \"\"\"\n",
    "        img = Image.open(self.images[index]).convert('RGB')\n",
    "        description = LADDDataSET.parse_voc_xml(\n",
    "            ET.parse(self.annotations[index]).getroot())\n",
    "\n",
    "        # get bounding box coordinates \n",
    "        num_objs = len(description['annotation']['object'])\n",
    "        boxes = []\n",
    "        for l in description['annotation']['object']:\n",
    "            bb = l['bndbox']\n",
    "            boxes.append([int(bb['xmin']), int(bb['ymin']), int(bb['xmax']), int(bb['ymax'])])\n",
    "\n",
    "        target = {}\n",
    "        target[\"boxes\"] = torch.as_tensor(boxes, dtype=torch.float32)         # there is only one class            \n",
    "        target[\"labels\"] = labels = torch.ones((num_objs,), dtype=torch.int64)\n",
    "        \n",
    "        if self.transforms is not None:\n",
    "            img, target = self.transforms(img, target)\n",
    "\n",
    "        return img, target\n",
    "    \n",
    "\n",
    "    def __len__(self) -> int:\n",
    "        return len(self.images)\n",
    "\n",
    "    @staticmethod\n",
    "    def parse_voc_xml(node: ET.Element) -> Dict[str, Any]:\n",
    "        voc_dict: Dict[str, Any] = {}\n",
    "        children = list(node)\n",
    "        if children:\n",
    "            def_dic: Dict[str, Any] = collections.defaultdict(list)\n",
    "            for dc in map(LADDDataSET.parse_voc_xml, children):\n",
    "                for ind, v in dc.items():\n",
    "                    def_dic[ind].append(v)\n",
    "            if node.tag == 'annotation':\n",
    "                def_dic['object'] = [def_dic['object']]\n",
    "            voc_dict = {\n",
    "                node.tag:\n",
    "                    {ind: v[0] if len(v) == 1 else v\n",
    "                     for ind, v in def_dic.items()}\n",
    "            }\n",
    "        if node.text:\n",
    "            text = node.text.strip()\n",
    "            if not children:\n",
    "                voc_dict[node.tag] = text\n",
    "        return voc_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total images 1020  non empty: 987\n"
     ]
    }
   ],
   "source": [
    "# Pytorch implemenation of retinanet doesn't supports train on Images without any objects (which, probably need to be fixed)\n",
    "# see https://github.com/pytorch/vision/blob/master/torchvision/models/detection/retinanet.py#L475\n",
    "# As a temporary solution, yet, we just filtering out empty images\n",
    "\n",
    "splits_dir = os.path.join(voc_root, 'ImageSets/Main') \n",
    "annotation_dir = os.path.join(voc_root, 'Annotations')\n",
    "\n",
    "with open(os.path.join(splits_dir,'train.txt'), \"r\") as f:\n",
    "    file_names = [x.strip() for x in f.readlines()]\n",
    "\n",
    "non_empty = []\n",
    "for a in file_names:\n",
    "    description = LADDDataSET.parse_voc_xml(\n",
    "        ET.parse(os.path.join(annotation_dir, a + \".xml\")).getroot()\n",
    "    )\n",
    "    num_objs = len(description['annotation']['object'])\n",
    "    if num_objs > 0:\n",
    "        non_empty.append(a+'\\n')\n",
    "        \n",
    "with open(os.path.join(splits_dir,'train_non_empty.txt'), \"w\") as f:\n",
    "    f.writelines(non_empty)\n",
    "\n",
    "print('Total images '+str(len(file_names)), ' non empty: '+str(len(non_empty)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test DS\n",
    "im_idx = 99\n",
    "\n",
    "dataset = LADDDataSET(voc_root,'test',get_transform(train=True,target_size=params['target_size'])) \n",
    "(image,target) = dataset[im_idx] \n",
    "im = F.to_pil_image(image)\n",
    "draw = ImageDraw.Draw(im)\n",
    "\n",
    "for bb in target['boxes']:\n",
    "    draw.line([(bb[0], bb[1]), (bb[0], bb[3]), (bb[2], bb[3]),\n",
    "               (bb[2], bb[1]), (bb[0], bb[1])], width=4, fill=(255, 0, 0))\n",
    "\n",
    "im.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_train = LADDDataSET(voc_root,'train_non_empty',get_transform(train=True,target_size=params['target_size'])) \n",
    "dataset_val = LADDDataSET(voc_root,'val',get_transform(train=False,target_size=params['target_size'])) \n",
    "\n",
    "# define training and validation data loaders\n",
    "data_loader = torch.utils.data.DataLoader(\n",
    "    dataset_train, batch_size=params['batch_size'], shuffle=True, num_workers=4\n",
    "     ,collate_fn=collate_fn\n",
    ")\n",
    "\n",
    "data_loader_val = torch.utils.data.DataLoader(\n",
    "    dataset_val, batch_size=1, shuffle=False, num_workers=4\n",
    "     ,collate_fn=collate_fn\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Downloading: \"https://download.pytorch.org/models/resnet50-19c8e357.pth\" to /home/gosha20777/.cache/torch/hub/checkpoints/resnet50-19c8e357.pth\n",
      "24.0%IOPub message rate exceeded.\n",
      "The notebook server will temporarily stop sending output\n",
      "to the client in order to avoid crashing it.\n",
      "To change this limit, set the config variable\n",
      "`--NotebookApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "NotebookApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "NotebookApp.rate_limit_window=3.0 (secs)\n",
      "\n",
      "60.4%IOPub message rate exceeded.\n",
      "The notebook server will temporarily stop sending output\n",
      "to the client in order to avoid crashing it.\n",
      "To change this limit, set the config variable\n",
      "`--NotebookApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "NotebookApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "NotebookApp.rate_limit_window=3.0 (secs)\n",
      "\n",
      "95.1%IOPub message rate exceeded.\n",
      "The notebook server will temporarily stop sending output\n",
      "to the client in order to avoid crashing it.\n",
      "To change this limit, set the config variable\n",
      "`--NotebookApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "NotebookApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "NotebookApp.rate_limit_window=3.0 (secs)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "model = torchvision.models.detection.fasterrcnn_resnet50_fpn(\n",
    "    pretrained=False,\n",
    "    num_classes=2,\n",
    "    pretrained_backbone=True, \n",
    "    min_size=params['target_size'][0],\n",
    "    max_size=params['target_size'][1],\n",
    "    trainable_backbone_layers = 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device('cuda')\n",
    "model = model.to(device)\n",
    "\n",
    "optimizer = torch.optim.SGD(\n",
    "    model.parameters(), \n",
    "    lr=params['lr'],\n",
    "    momentum=0.9,\n",
    "    weight_decay=0.0005) \n",
    "lr_scheduler = torch.optim.lr_scheduler.StepLR(\n",
    "    optimizer,\n",
    "    step_size=3,\n",
    "    gamma=0.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "ename": "RuntimeError",
     "evalue": "CUDA out of memory. Tried to allocate 186.00 MiB (GPU 0; 3.82 GiB total capacity; 1.90 GiB already allocated; 197.06 MiB free; 2.10 GiB reserved in total by PyTorch)",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-22-6c0d5c0fa52b>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mepoch\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m2\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m     \u001b[0mtrain_one_epoch\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moptimizer\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdata_loader\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdevice\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mepoch\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mprint_freq\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m100\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      3\u001b[0m     \u001b[0mprint\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0;34m\"Train done, evaluating.\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m     \u001b[0;31m# update the learning rate\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m     \u001b[0mlr_scheduler\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstep\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-4-0555bc5ea552>\u001b[0m in \u001b[0;36mtrain_one_epoch\u001b[0;34m(model, optimizer, data_loader, device, epoch, print_freq)\u001b[0m\n\u001b[1;32m    224\u001b[0m         \u001b[0mtargets\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m{\u001b[0m\u001b[0mk\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mv\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mk\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mv\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mt\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mitems\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m}\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mt\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mtargets\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    225\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 226\u001b[0;31m         \u001b[0mloss_dict\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mimages\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtargets\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    227\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    228\u001b[0m         \u001b[0mlosses\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0msum\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mloss\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mloss\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mloss_dict\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/pytorch-baseline/lib/python3.8/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m    887\u001b[0m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    888\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 889\u001b[0;31m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    890\u001b[0m         for hook in itertools.chain(\n\u001b[1;32m    891\u001b[0m                 \u001b[0m_global_forward_hooks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/pytorch-baseline/lib/python3.8/site-packages/torchvision/models/detection/generalized_rcnn.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, images, targets)\u001b[0m\n\u001b[1;32m     92\u001b[0m                                      .format(degen_bb, target_idx))\n\u001b[1;32m     93\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 94\u001b[0;31m         \u001b[0mfeatures\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackbone\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mimages\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtensors\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     95\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfeatures\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mTensor\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     96\u001b[0m             \u001b[0mfeatures\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mOrderedDict\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'0'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfeatures\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/pytorch-baseline/lib/python3.8/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m    887\u001b[0m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    888\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 889\u001b[0;31m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    890\u001b[0m         for hook in itertools.chain(\n\u001b[1;32m    891\u001b[0m                 \u001b[0m_global_forward_hooks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/pytorch-baseline/lib/python3.8/site-packages/torchvision/models/detection/backbone_utils.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m     42\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     43\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 44\u001b[0;31m         \u001b[0mx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbody\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     45\u001b[0m         \u001b[0mx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfpn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     46\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/pytorch-baseline/lib/python3.8/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m    887\u001b[0m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    888\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 889\u001b[0;31m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    890\u001b[0m         for hook in itertools.chain(\n\u001b[1;32m    891\u001b[0m                 \u001b[0m_global_forward_hooks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/pytorch-baseline/lib/python3.8/site-packages/torchvision/models/_utils.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m     61\u001b[0m         \u001b[0mout\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mOrderedDict\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     62\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mname\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmodule\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mitems\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 63\u001b[0;31m             \u001b[0mx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodule\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     64\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mname\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mreturn_layers\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     65\u001b[0m                 \u001b[0mout_name\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mreturn_layers\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mname\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/pytorch-baseline/lib/python3.8/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m    887\u001b[0m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    888\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 889\u001b[0;31m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    890\u001b[0m         for hook in itertools.chain(\n\u001b[1;32m    891\u001b[0m                 \u001b[0m_global_forward_hooks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/pytorch-baseline/lib/python3.8/site-packages/torchvision/ops/misc.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m     94\u001b[0m         \u001b[0mscale\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mw\u001b[0m \u001b[0;34m*\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mrv\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0meps\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrsqrt\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     95\u001b[0m         \u001b[0mbias\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mb\u001b[0m \u001b[0;34m-\u001b[0m \u001b[0mrm\u001b[0m \u001b[0;34m*\u001b[0m \u001b[0mscale\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 96\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mx\u001b[0m \u001b[0;34m*\u001b[0m \u001b[0mscale\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mbias\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     97\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     98\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m__repr__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0mstr\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mRuntimeError\u001b[0m: CUDA out of memory. Tried to allocate 186.00 MiB (GPU 0; 3.82 GiB total capacity; 1.90 GiB already allocated; 197.06 MiB free; 2.10 GiB reserved in total by PyTorch)"
     ]
    }
   ],
   "source": [
    "for epoch in range(2):\n",
    "    train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq=100)\n",
    "    print (\"Train done, evaluating.\")\n",
    "    # update the learning rate\n",
    "    lr_scheduler.step()\n",
    "    # evaluate on the test dataset\n",
    "    inference_res = evaluate(model,data_loader_val)\n",
    "    print('Inference done, computing mAp : ')\n",
    "    print(evaluate_res(inference_res, iou_threshold = 0.5, score_threshold = 0.05))    \n",
    "    print(evaluate_res(inference_res, iou_threshold = 0.6, score_threshold = 0.05))\n",
    "    print('Epoch Done')\n",
    "    \n",
    "torch.save(model.state_dict(), 'weights/resnet50_FRCNN_baseline.pth')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_test = LADDDataSET(voc_root,'test',get_transform(train=False,target_size=params['target_size'])) \n",
    "data_loader_test = torch.utils.data.DataLoader(\n",
    "    dataset_val, batch_size=1, shuffle=False, num_workers=1\n",
    "     ,collate_fn=collate_fn\n",
    ")\n",
    "\n",
    "image_idx = 0\n",
    "\n",
    "cpu_device = torch.device(\"cpu\")\n",
    "model.eval()\n",
    "for images, targets in data_loader_test:\n",
    "    g_images = list(img.to(device) for img in images)\n",
    "\n",
    "    if torch.cuda.is_available():\n",
    "        torch.cuda.synchronize()\n",
    "    outputs = model(g_images)\n",
    "\n",
    "    outputs = [{k: v.to(cpu_device) for k, v in t.items()} for t in outputs]\n",
    "    res = targets, outputs\n",
    "    break\n",
    "\n",
    "\n",
    "im = F.to_pil_image(images[image_idx])\n",
    "targets\n",
    "# im = to_pil_image(dataset[10][0])\n",
    "draw = ImageDraw.Draw(im)\n",
    "\n",
    "for idx in range(len(outputs[image_idx]['boxes'])):\n",
    "    width = math.ceil(outputs[image_idx]['scores'][idx]*10)\n",
    "    bb = outputs[0]['boxes'][idx]\n",
    "    draw.line([(bb[0], bb[1]), (bb[0], bb[3]), (bb[2], bb[3]),\n",
    "               (bb[2], bb[1]), (bb[0], bb[1])], width=width, fill=(255, 0, 0))\n",
    "\n",
    "for bb in targets[image_idx]['boxes'][:10]:\n",
    "    draw.line([(bb[0], bb[1]), (bb[0], bb[3]), (bb[2], bb[3]),\n",
    "               (bb[2], bb[1]), (bb[0], bb[1])], width=4, fill=(0,255, 0))\n",
    "im.show()"
   ]
  },
  {
   "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}