Yolov1.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "#We use VOC2012 dataset, you should download it, and put it in ../data.\n",
    "#And you also need create a folder called 'labels'\n",
    "DATA_PATH = \"../data/VOC2012/\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "#First, we should transform the dataset.\n",
    "#We need the data type of (cls, x, y, w, h)\n",
    "CLASSES = ['person', 'bird', 'cat', 'cow', 'dog', 'horse', 'sheep',\n",
    "           'aeroplane', 'bicycle', 'boat', 'bus', 'car', 'motorbike', 'train',\n",
    "           'bottle', 'chair', 'dining table', 'potted plant', 'sofa', 'tvmonitor']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import xml.etree.ElementTree as ET\n",
    "import os\n",
    "import cv2\n",
    "\n",
    "def convert(size, box):\n",
    "    dw = 1. / size[0]\n",
    "    dh = 1. / size[1]\n",
    "    \n",
    "    x = (box[0] + box[1]) / 2.0\n",
    "    y = (box[2] + box[3]) / 2.0\n",
    "    w = box[1] - box[0]\n",
    "    h = box[3] - box[2]\n",
    "    \n",
    "    x = x * dw\n",
    "    y = y * dh\n",
    "    w = w * dw\n",
    "    h = h * dh\n",
    "    \n",
    "    return (x, y, w ,h)\n",
    "\n",
    "def convert_annotation(image_id):\n",
    "    in_file = open(DATA_PATH + 'Annotations/%s' % image_id)\n",
    "    image_id = image_id.split('.')[0]\n",
    "    out_file = open('./labels/%s.txt' % image_id, 'w')\n",
    "    \n",
    "    tree = ET.parse(in_file)\n",
    "    root = tree.getroot()\n",
    "    size = root.find('size')\n",
    "    \n",
    "    w = int(size.find('width').text)\n",
    "    h = int(size.find('height').text)\n",
    "    \n",
    "    for obj in root.iter('object'):\n",
    "        difficult = obj.find('difficult').text\n",
    "        cls = obj.find('name').text\n",
    "        \n",
    "        if cls not in CLASSES or int(difficult) == 1:\n",
    "            continue\n",
    "        \n",
    "        cls_id = CLASSES.index(cls)\n",
    "        xmlbox = obj.find('bndbox')\n",
    "        points = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text),\n",
    "             float(xmlbox.find('ymax').text))\n",
    "        convert_value = convert((w, h), points)\n",
    "        \n",
    "        out_file.write(str(cls_id) + \" \" + \" \".join([str(a) for a in convert_value]) + '\\n')\n",
    "\n",
    "def make_label_txt():\n",
    "    filenames = os.listdir(DATA_PATH + 'Annotations')\n",
    "    for file in filenames:\n",
    "        convert_annotation(file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Checkout the convert function\n",
    "def show_labels_img(img_name):\n",
    "    img = cv2.imread(DATA_PATH + 'JPEGImages/' + img_name + '.jpg')\n",
    "    h, w = img.shape[:2]\n",
    "    label = []\n",
    "    \n",
    "    with open('./labels/' + img_name + '.txt', 'r') as labels:\n",
    "        for label in labels:\n",
    "            label = label.split(' ')\n",
    "            label = [float(x.strip()) for x in label]\n",
    "            print(CLASSES[int(label[0])])\n",
    "            pt1 = (int(label[1] * w - label[3] * w / 2), int(label[2] * h - label[4] * h / 2))\n",
    "            pt2 = (int(label[1] * w + label[3] * w / 2), int(label[2] * h + label[4] * h / 2))\n",
    "            cv2.putText(img, CLASSES[int(label[0])], pt1, cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,0,255))\n",
    "            cv2.rectangle(img,pt1,pt2,(0,0,255,2))\n",
    "    \n",
    "    cv2.imshow('img', img)\n",
    "    cv2.waitKey(0)\n",
    "\n",
    "make_label_txt()\n",
    "show_labels_img('2007_000027')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Prepare Pytorch Dataset.\n",
    "import numpy as np\n",
    "import torch\n",
    "import torchvision.transforms as transforms\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import torch.nn as nn\n",
    "import torchvision.models as tvmodel\n",
    "import visdom as visdom\n",
    "\n",
    "NUM_BBOX = 2\n",
    "\n",
    "class VOC2012(Dataset):\n",
    "    def __init__(self, is_train=True, is_aug=True):\n",
    "        self.filenames = []\n",
    "        \n",
    "        if is_train:\n",
    "            with open(DATA_PATH + 'ImageSets/Main/train.txt', 'r') as files:\n",
    "                self.filenames = [file.strip() for file in files]\n",
    "        else:\n",
    "            with open(DATA_PATH + 'ImageSets/Main/val.txt', 'r') as files:\n",
    "                self.filenames = [file.strip() for file in files]\n",
    "        \n",
    "        self.img_path = DATA_PATH + 'JPEGImages/'\n",
    "        self.label_path = './labels/'\n",
    "        self.is_aug = is_aug\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.filenames)\n",
    "    \n",
    "    def __getitem__(self, item):\n",
    "        img = cv2.imread(self.img_path + self.filenames[item] + '.jpg')\n",
    "        h, w = img.shape[0: 2]\n",
    "        input_size = 448\n",
    "        \n",
    "        padding_w, padding_h = 0, 0\n",
    "        \n",
    "        if h > w:\n",
    "            padding_w = (h - w) // 2\n",
    "            img = np.pad(img, ((0, 0), (padding_w, padding_w), (0, 0)), 'constant', constant_values=0)\n",
    "        elif w > h:\n",
    "            padding_h = (w - h) // 2\n",
    "            img = np.pad(img, ((padding_h, padding_h), (0, 0), (0, 0)), 'constant', constant_values=0)\n",
    "        \n",
    "        img = cv2.resize(img, (input_size, input_size))\n",
    "        \n",
    "        if self.is_aug:\n",
    "            aug = transforms.Compose([\n",
    "                transforms.ToTensor()\n",
    "            ])\n",
    "            img = aug(img)\n",
    "\n",
    "        with open(self.label_path + self.filenames[item] + \".txt\") as file:\n",
    "            bbox = file.read().split('\\n')\n",
    "        bbox = [x.split() for x in bbox]\n",
    "        bbox = [float(x) for y in bbox for x in y]\n",
    "        if len(bbox) % 5 != 0:\n",
    "            raise ValueError(\"File:\" + self.label_path + self.filenames[item] + \".txt\" + \"——bbox Extraction Error!\")\n",
    "\n",
    "        for i in range(len(bbox) // 5):\n",
    "            if padding_w != 0:\n",
    "                bbox[i * 5 + 1] = (bbox[i * 5 + 1] * w + padding_w) / h\n",
    "                bbox[i * 5 + 3] = (bbox[i * 5 + 3] * w) / h\n",
    "            elif padding_h != 0:\n",
    "                bbox[i * 5 + 2] = (bbox[i * 5 + 2] * h + padding_h) / w\n",
    "                bbox[i * 5 + 4] = (bbox[i * 5 + 4] * h) / w\n",
    "\n",
    "        labels = convert_bbox2labels(bbox)\n",
    "        labels = transforms.ToTensor()(labels)\n",
    "        \n",
    "        return img, labels\n",
    "\n",
    "def convert_bbox2labels(bbox):\n",
    "    \n",
    "    gridsize = 1.0 / 7\n",
    "    labels = np.zeros((7, 7, 5 * NUM_BBOX + len(CLASSES)))  # 注意，此处需要根据不同数据集的类别个数进行修改\n",
    "    for i in range(len(bbox) // 5):\n",
    "        gridx = int(bbox[i * 5 + 1] // gridsize)  # 当前bbox中心落在第gridx个网格,列\n",
    "        gridy = int(bbox[i * 5 + 2] // gridsize)  # 当前bbox中心落在第gridy个网格,行\n",
    "        # (bbox中心坐标 - 网格左上角点的坐标)/网格大小  ==> bbox中心点的相对位置\n",
    "        gridpx = bbox[i * 5 + 1] / gridsize - gridx\n",
    "        gridpy = bbox[i * 5 + 2] / gridsize - gridy\n",
    "        # 将第gridy行，gridx列的网格设置为负责当前ground truth的预测，置信度和对应类别概率均置为1\n",
    "        labels[gridy, gridx, 0:5] = np.array([gridpx, gridpy, bbox[i * 5 + 3], bbox[i * 5 + 4], 1])\n",
    "        labels[gridy, gridx, 5:10] = np.array([gridpx, gridpy, bbox[i * 5 + 3], bbox[i * 5 + 4], 1])\n",
    "        labels[gridy, gridx, 10+int(bbox[i*5])] = 1\n",
    "    return labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "#define loss function\n",
    "class Loss_yolov1(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Loss_yolov1, self).__init__()\n",
    "    \n",
    "    def forward(self, pred, labels):\n",
    "        num_gridx, num_gridy = labels.size()[-2:]  # 划分网格数量\n",
    "        num_b = 2  # 每个网格的bbox数量\n",
    "        num_cls = 20  # 类别数量\n",
    "        noobj_confi_loss = 0.  # 不含目标的网格损失(只有置信度损失)\n",
    "        coor_loss = 0.  # 含有目标的bbox的坐标损失\n",
    "        obj_confi_loss = 0.  # 含有目标的bbox的置信度损失\n",
    "        class_loss = 0.  # 含有目标的网格的类别损失\n",
    "        n_batch = labels.size()[0]  # batchsize的大小\n",
    "\n",
    "        # 可以考虑用矩阵运算进行优化，提高速度，为了准确起见，这里还是用循环\n",
    "        for i in range(n_batch):  # batchsize循环\n",
    "            for n in range(7):  # x方向网格循环\n",
    "                for m in range(7):  # y方向网格循环\n",
    "                    if labels[i,4,m,n]==1:# 如果包含物体\n",
    "                        # 将数据(px,py,w,h)转换为(x1,y1,x2,y2)\n",
    "                        # 先将px,py转换为cx,cy，即相对网格的位置转换为标准化后实际的bbox中心位置cx,xy\n",
    "                        # 然后再利用(cx-w/2,cy-h/2,cx+w/2,cy+h/2)转换为xyxy形式，用于计算iou\n",
    "                        bbox1_pred_xyxy = ((pred[i,0,m,n]+m)/num_gridx - pred[i,2,m,n]/2,(pred[i,1,m,n]+n)/num_gridy - pred[i,3,m,n]/2,\n",
    "                                           (pred[i,0,m,n]+m)/num_gridx + pred[i,2,m,n]/2,(pred[i,1,m,n]+n)/num_gridy + pred[i,3,m,n]/2)\n",
    "                        bbox2_pred_xyxy = ((pred[i,5,m,n]+m)/num_gridx - pred[i,7,m,n]/2,(pred[i,6,m,n]+n)/num_gridy - pred[i,8,m,n]/2,\n",
    "                                           (pred[i,5,m,n]+m)/num_gridx + pred[i,7,m,n]/2,(pred[i,6,m,n]+n)/num_gridy + pred[i,8,m,n]/2)\n",
    "                        bbox_gt_xyxy = ((labels[i,0,m,n]+m)/num_gridx - labels[i,2,m,n]/2,(labels[i,1,m,n]+n)/num_gridy - labels[i,3,m,n]/2,\n",
    "                                        (labels[i,0,m,n]+m)/num_gridx + labels[i,2,m,n]/2,(labels[i,1,m,n]+n)/num_gridy + labels[i,3,m,n]/2)\n",
    "                        iou1 = calculate_iou(bbox1_pred_xyxy,bbox_gt_xyxy)\n",
    "                        iou2 = calculate_iou(bbox2_pred_xyxy,bbox_gt_xyxy)\n",
    "                        # 选择iou大的bbox作为负责物体\n",
    "                        if iou1 >= iou2:\n",
    "                            coor_loss = coor_loss + 5 * (torch.sum((pred[i,0:2,m,n] - labels[i,0:2,m,n])**2) \\\n",
    "                                        + torch.sum((pred[i,2:4,m,n].sqrt()-labels[i,2:4,m,n].sqrt())**2))\n",
    "                            obj_confi_loss = obj_confi_loss + (pred[i,4,m,n] - iou1)**2\n",
    "                            # iou比较小的bbox不负责预测物体，因此confidence loss算在noobj中，注意，对于标签的置信度应该是iou2\n",
    "                            noobj_confi_loss = noobj_confi_loss + 0.5 * ((pred[i,9,m,n]-iou2)**2)\n",
    "                        else:\n",
    "                            coor_loss = coor_loss + 5 * (torch.sum((pred[i,5:7,m,n] - labels[i,5:7,m,n])**2) \\\n",
    "                                        + torch.sum((pred[i,7:9,m,n].sqrt()-labels[i,7:9,m,n].sqrt())**2))\n",
    "                            obj_confi_loss = obj_confi_loss + (pred[i,9,m,n] - iou2)**2\n",
    "                            # iou比较小的bbox不负责预测物体，因此confidence loss算在noobj中,注意，对于标签的置信度应该是iou1\n",
    "                            noobj_confi_loss = noobj_confi_loss + 0.5 * ((pred[i, 4, m, n]-iou1) ** 2)\n",
    "                        class_loss = class_loss + torch.sum((pred[i,10:,m,n] - labels[i,10:,m,n])**2)\n",
    "                    else:  # 如果不包含物体\n",
    "                        noobj_confi_loss = noobj_confi_loss + 0.5 * torch.sum(pred[i,[4,9],m,n]**2)\n",
    "\n",
    "        loss = coor_loss + obj_confi_loss + noobj_confi_loss + class_loss\n",
    "        # 此处可以写代码验证一下loss的大致计算是否正确，这个要验证起来比较麻烦，比较简洁的办法是，将输入的pred置为全1矩阵，再进行误差检查，会直观很多。\n",
    "        return loss/n_batch\n",
    "\n",
    "def calculate_iou(bbox1,bbox2):\n",
    "    \"\"\"计算bbox1=(x1,y1,x2,y2)和bbox2=(x3,y3,x4,y4)两个bbox的iou\"\"\"\n",
    "    intersect_bbox = [0., 0., 0., 0.]  # bbox1和bbox2的交集\n",
    "    if bbox1[2]<bbox2[0] or bbox1[0]>bbox2[2] or bbox1[3]<bbox2[1] or bbox1[1]>bbox2[3]:\n",
    "        pass\n",
    "    else:\n",
    "        intersect_bbox[0] = max(bbox1[0],bbox2[0])\n",
    "        intersect_bbox[1] = max(bbox1[1],bbox2[1])\n",
    "        intersect_bbox[2] = min(bbox1[2],bbox2[2])\n",
    "        intersect_bbox[3] = min(bbox1[3],bbox2[3])\n",
    "\n",
    "    area1 = (bbox1[2] - bbox1[0]) * (bbox1[3] - bbox1[1])  # bbox1面积\n",
    "    area2 = (bbox2[2] - bbox2[0]) * (bbox2[3] - bbox2[1])  # bbox2面积\n",
    "    area_intersect = (intersect_bbox[2] - intersect_bbox[0]) * (intersect_bbox[3] - intersect_bbox[1])  # 交集面积\n",
    "    # print(bbox1,bbox2)\n",
    "    # print(intersect_bbox)\n",
    "    # input()\n",
    "\n",
    "    if area_intersect>0:\n",
    "        return area_intersect / (area1 + area2 - area_intersect)  # 计算iou\n",
    "    else:\n",
    "        return 0\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "#define model\n",
    "class Yolov1_resnet(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Yolov1_resnet, self).__init__()\n",
    "        resnet = tvmodel.resnet34(pretrained=True)\n",
    "        resnet_out_channel = resnet.fc.in_features\n",
    "        self.resnet = nn.Sequential(*list(resnet.children())[: -2])\n",
    "        \n",
    "        self.Conv_layers = nn.Sequential(\n",
    "            nn.Conv2d(resnet_out_channel, 1024, 3, padding=1),\n",
    "            nn.BatchNorm2d(1024),\n",
    "            nn.LeakyReLU(),\n",
    "            \n",
    "            nn.Conv2d(1024, 1024, 3, stride=2, padding=1),\n",
    "            nn.BatchNorm2d(1024),\n",
    "            nn.LeakyReLU(),\n",
    "            \n",
    "            nn.Conv2d(1024, 1024, 3, padding=1),\n",
    "            nn.BatchNorm2d(1024),\n",
    "            nn.LeakyReLU(),\n",
    "            \n",
    "            nn.Conv2d(1024, 1024, 3, padding=1),\n",
    "            nn.BatchNorm2d(1024),\n",
    "            nn.LeakyReLU(),\n",
    "        )\n",
    "        \n",
    "        self.Fc_layers = nn.Sequential(\n",
    "            nn.Linear(7 * 7 * 1024, 4096),\n",
    "            nn.LeakyReLU(),\n",
    "            \n",
    "            nn.Linear(4096, 7 * 7 * 30),\n",
    "            nn.Sigmoid()\n",
    "        )\n",
    "        \n",
    "    def forward(self, out):\n",
    "        out = self.resnet(out)\n",
    "        out = self.Conv_layers(out)\n",
    "        out = out.view(out.size()[0], -1)\n",
    "        out = self.Fc_layers(out)\n",
    "        return out.reshape(-1, (5 * NUM_BBOX + len(CLASSES)), 7, 7)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Yolov1_resnet(\n",
      "  (resnet): Sequential(\n",
      "    (0): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)\n",
      "    (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    (2): ReLU(inplace=True)\n",
      "    (3): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)\n",
      "    (4): Sequential(\n",
      "      (0): BasicBlock(\n",
      "        (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (relu): ReLU(inplace=True)\n",
      "        (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "      (1): BasicBlock(\n",
      "        (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (relu): ReLU(inplace=True)\n",
      "        (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "      (2): BasicBlock(\n",
      "        (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (relu): ReLU(inplace=True)\n",
      "        (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "    )\n",
      "    (5): Sequential(\n",
      "      (0): BasicBlock(\n",
      "        (conv1): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
      "        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (relu): ReLU(inplace=True)\n",
      "        (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (downsample): Sequential(\n",
      "          (0): Conv2d(64, 128, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
      "          (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        )\n",
      "      )\n",
      "      (1): BasicBlock(\n",
      "        (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (relu): ReLU(inplace=True)\n",
      "        (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "      (2): BasicBlock(\n",
      "        (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (relu): ReLU(inplace=True)\n",
      "        (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "      (3): BasicBlock(\n",
      "        (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (relu): ReLU(inplace=True)\n",
      "        (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "    )\n",
      "    (6): Sequential(\n",
      "      (0): BasicBlock(\n",
      "        (conv1): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
      "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (relu): ReLU(inplace=True)\n",
      "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (downsample): Sequential(\n",
      "          (0): Conv2d(128, 256, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
      "          (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        )\n",
      "      )\n",
      "      (1): BasicBlock(\n",
      "        (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (relu): ReLU(inplace=True)\n",
      "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "      (2): BasicBlock(\n",
      "        (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (relu): ReLU(inplace=True)\n",
      "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "      (3): BasicBlock(\n",
      "        (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (relu): ReLU(inplace=True)\n",
      "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "      (4): BasicBlock(\n",
      "        (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (relu): ReLU(inplace=True)\n",
      "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "      (5): BasicBlock(\n",
      "        (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (relu): ReLU(inplace=True)\n",
      "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "    )\n",
      "    (7): Sequential(\n",
      "      (0): BasicBlock(\n",
      "        (conv1): Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
      "        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (relu): ReLU(inplace=True)\n",
      "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (downsample): Sequential(\n",
      "          (0): Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
      "          (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        )\n",
      "      )\n",
      "      (1): BasicBlock(\n",
      "        (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (relu): ReLU(inplace=True)\n",
      "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "      (2): BasicBlock(\n",
      "        (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (relu): ReLU(inplace=True)\n",
      "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "    )\n",
      "  )\n",
      "  (Conv_layers): Sequential(\n",
      "    (0): Conv2d(512, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    (1): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    (2): LeakyReLU(negative_slope=0.01)\n",
      "    (3): Conv2d(1024, 1024, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "    (4): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    (5): LeakyReLU(negative_slope=0.01)\n",
      "    (6): Conv2d(1024, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    (7): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    (8): LeakyReLU(negative_slope=0.01)\n",
      "    (9): Conv2d(1024, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    (10): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    (11): LeakyReLU(negative_slope=0.01)\n",
      "  )\n",
      "  (Fc_layers): Sequential(\n",
      "    (0): Linear(in_features=50176, out_features=4096, bias=True)\n",
      "    (1): LeakyReLU(negative_slope=0.01)\n",
      "    (2): Linear(in_features=4096, out_features=1470, bias=True)\n",
      "    (3): Sigmoid()\n",
      "  )\n",
      ")\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 30, 7, 7])\n"
     ]
    }
   ],
   "source": [
    "#test model\n",
    "x = torch.randn((1, 3, 448, 448))\n",
    "net = Yolov1_resnet()\n",
    "print(net)\n",
    "y = net(x)\n",
    "print(y.size())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "ename": "RuntimeError",
     "evalue": "CUDA out of memory. Tried to allocate 1.20 GiB (GPU 0; 4.00 GiB total capacity; 2.41 GiB already allocated; 604.14 MiB free; 2.44 GiB reserved in total by PyTorch)",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "\u001b[1;32m<ipython-input-8-62b9a56f85ab>\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[0;32m     29\u001b[0m         \u001b[0minputs\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0minputs\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mto\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mdevice\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     30\u001b[0m         \u001b[0mlabels\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mlabels\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mfloat\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mto\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mdevice\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m---> 31\u001b[1;33m         \u001b[0mpred\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mmodel\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0minputs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m     32\u001b[0m         \u001b[0mloss\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mcriterion\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mpred\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mlabels\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     33\u001b[0m         \u001b[0moptimizer\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mzero_grad\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32md:\\Anaconda3\\envs\\pytorch1.5\\lib\\site-packages\\torch\\nn\\modules\\module.py\u001b[0m in \u001b[0;36m__call__\u001b[1;34m(self, *input, **kwargs)\u001b[0m\n\u001b[0;32m    548\u001b[0m             \u001b[0mresult\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    549\u001b[0m         \u001b[1;32melse\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 550\u001b[1;33m             \u001b[0mresult\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    551\u001b[0m         \u001b[1;32mfor\u001b[0m \u001b[0mhook\u001b[0m \u001b[1;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_forward_hooks\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    552\u001b[0m             \u001b[0mhook_result\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mhook\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m<ipython-input-7-94882f87beb3>\u001b[0m in \u001b[0;36mforward\u001b[1;34m(self, out)\u001b[0m\n\u001b[0;32m     34\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     35\u001b[0m     \u001b[1;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mout\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m---> 36\u001b[1;33m         \u001b[0mout\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mresnet\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mout\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m     37\u001b[0m         \u001b[0mout\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mConv_layers\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mout\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     38\u001b[0m         \u001b[0mout\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mout\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mview\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mout\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0msize\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m[\u001b[0m\u001b[1;36m0\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m-\u001b[0m\u001b[1;36m1\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32md:\\Anaconda3\\envs\\pytorch1.5\\lib\\site-packages\\torch\\nn\\modules\\module.py\u001b[0m in \u001b[0;36m__call__\u001b[1;34m(self, *input, **kwargs)\u001b[0m\n\u001b[0;32m    548\u001b[0m             \u001b[0mresult\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    549\u001b[0m         \u001b[1;32melse\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 550\u001b[1;33m             \u001b[0mresult\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    551\u001b[0m         \u001b[1;32mfor\u001b[0m \u001b[0mhook\u001b[0m \u001b[1;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_forward_hooks\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    552\u001b[0m             \u001b[0mhook_result\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mhook\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32md:\\Anaconda3\\envs\\pytorch1.5\\lib\\site-packages\\torch\\nn\\modules\\container.py\u001b[0m in \u001b[0;36mforward\u001b[1;34m(self, input)\u001b[0m\n\u001b[0;32m     98\u001b[0m     \u001b[1;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0minput\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     99\u001b[0m         \u001b[1;32mfor\u001b[0m \u001b[0mmodule\u001b[0m \u001b[1;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 100\u001b[1;33m             \u001b[0minput\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mmodule\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0minput\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    101\u001b[0m         \u001b[1;32mreturn\u001b[0m \u001b[0minput\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    102\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32md:\\Anaconda3\\envs\\pytorch1.5\\lib\\site-packages\\torch\\nn\\modules\\module.py\u001b[0m in \u001b[0;36m__call__\u001b[1;34m(self, *input, **kwargs)\u001b[0m\n\u001b[0;32m    548\u001b[0m             \u001b[0mresult\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    549\u001b[0m         \u001b[1;32melse\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 550\u001b[1;33m             \u001b[0mresult\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    551\u001b[0m         \u001b[1;32mfor\u001b[0m \u001b[0mhook\u001b[0m \u001b[1;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_forward_hooks\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    552\u001b[0m             \u001b[0mhook_result\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mhook\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32md:\\Anaconda3\\envs\\pytorch1.5\\lib\\site-packages\\torch\\nn\\modules\\batchnorm.py\u001b[0m in \u001b[0;36mforward\u001b[1;34m(self, input)\u001b[0m\n\u001b[0;32m    104\u001b[0m             \u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mrunning_mean\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mrunning_var\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mweight\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mbias\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    105\u001b[0m             \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mtraining\u001b[0m \u001b[1;32mor\u001b[0m \u001b[1;32mnot\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mtrack_running_stats\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 106\u001b[1;33m             exponential_average_factor, self.eps)\n\u001b[0m\u001b[0;32m    107\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    108\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32md:\\Anaconda3\\envs\\pytorch1.5\\lib\\site-packages\\torch\\nn\\functional.py\u001b[0m in \u001b[0;36mbatch_norm\u001b[1;34m(input, running_mean, running_var, weight, bias, training, momentum, eps)\u001b[0m\n\u001b[0;32m   1921\u001b[0m     return torch.batch_norm(\n\u001b[0;32m   1922\u001b[0m         \u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mweight\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mbias\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mrunning_mean\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mrunning_var\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m-> 1923\u001b[1;33m         \u001b[0mtraining\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mmomentum\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0meps\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mtorch\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mbackends\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mcudnn\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0menabled\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m   1924\u001b[0m     )\n\u001b[0;32m   1925\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;31mRuntimeError\u001b[0m: CUDA out of memory. Tried to allocate 1.20 GiB (GPU 0; 4.00 GiB total capacity; 2.41 GiB already allocated; 604.14 MiB free; 2.44 GiB reserved in total by PyTorch)"
     ]
    }
   ],
   "source": [
    "#start training\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "epoch = 50\n",
    "batchsize = 100\n",
    "lr = 0.01\n",
    "\n",
    "train_data = VOC2012()\n",
    "train_dataloader = DataLoader(VOC2012(is_train=True), batch_size=batchsize, shuffle=True)\n",
    "\n",
    "model = Yolov1_resnet().to(device)\n",
    "\n",
    "for layer in model.children():\n",
    "    layer.requires_grad = False\n",
    "    break\n",
    "criterion = Loss_yolov1()\n",
    "optimizer = torch.optim.SGD(model.parameters(), lr=lr, momentum=0.9, weight_decay=0.0005)\n",
    "\n",
    "is_vis = False\n",
    "\n",
    "if is_vis:\n",
    "    vis = visdom.Visdom()\n",
    "    viswin1 = vis.line(np.array([0.]),np.array([0.]),opts=dict(title=\"Loss/Step\",xlabel=\"100*step\",ylabel=\"Loss\"))\n",
    "\n",
    "for e in range(epoch):\n",
    "    model.train()\n",
    "    yl = torch.Tensor([0]).to(device)\n",
    "    for i, (inputs, labels) in enumerate(train_dataloader):\n",
    "        inputs = inputs.to(device)\n",
    "        labels = labels.float().to(device)\n",
    "        pred = model(inputs)\n",
    "        loss = criterion(pred, labels)\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        \n",
    "        print(\"Epoch %d/%d| Step %d/%d| Loss: %.2f\"%(e,epoch,i,len(train_data)//batchsize,loss))\n",
    "        yl = yl + loss\n",
    "        if is_vis and (i + 1) % 100 == 0:\n",
    "            vis.line(np.array([yl.cpu().item()/(i+1)]),np.array([i+e*len(train_data)//batchsize]),win=viswin1,update='append')\n",
    "    if (e + 1) % 10 == 0:\n",
    "        torch.save(model, \"./models_pkl/YOLOv1_epoch\" + str(e + 1) + \".pkl\")\n",
    "        # compute_val_map(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Pytorch for Deeplearning",
   "language": "python",
   "name": "pytorch"
  },
  "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}