yolo_v1.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.functional as F"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "outputs": [
    {
     "data": {
      "text/plain": "torch.Size([1, 3, 448, 448])"
     },
     "execution_count": 78,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "xin = torch.randn((1, 3, 448, 448))\n",
    "xin.shape"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 168,
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.functional as F\n",
    "\n",
    "def repeat_block(blocks: list[torch.nn.Module], repeats: int) -> list[torch.nn.Module]:\n",
    "    modules = [block for i in range(repeats) for block in blocks]\n",
    "    return modules\n",
    "\n",
    "\n",
    "class YoloV1(torch.nn.Module):\n",
    "\n",
    "    def __init__(self, num_boxes: int = 2, n_classes:int = 20, split_size:int = 7, final_layer_size: int = 4096):\n",
    "        super().__init__()\n",
    "\n",
    "        self.C = n_classes\n",
    "        self.S = split_size\n",
    "        self.B = num_boxes\n",
    "        self.final_layer_size = final_layer_size\n",
    "\n",
    "        self.block_1 = torch.nn.Sequential(\n",
    "            ConvBlock(3, 64, kernel_size=(7, 7), stride=2, padding=3),\n",
    "            torch.nn.MaxPool2d(kernel_size=(2, 2), stride=2)\n",
    "        )\n",
    "\n",
    "        self.block_2 = torch.nn.Sequential(\n",
    "            ConvBlock(64, 192, kernel_size=(3, 3), padding=1),\n",
    "            torch.nn.MaxPool2d(kernel_size=(2, 2))\n",
    "        )\n",
    "\n",
    "        self.block_3 = torch.nn.Sequential(\n",
    "            ConvBlock(192, 128, kernel_size=(1, 1), padding=2),\n",
    "            ConvBlock(128, 256, kernel_size=(3, 3)),\n",
    "            ConvBlock(256, 256, kernel_size=(1, 1)),\n",
    "            ConvBlock(256, 512, kernel_size=(3, 3)),\n",
    "            torch.nn.MaxPool2d(kernel_size=(2, 2))\n",
    "        )\n",
    "\n",
    "        # BLOCK 4\n",
    "        self.block_4 = torch.nn.Sequential(\n",
    "            *repeat_block([ConvBlock(512, 256, kernel_size=(1, 1)), ConvBlock(256, 512, kernel_size=(3, 3), padding=1)], 4),\n",
    "            ConvBlock(512, 512, kernel_size=(1, 1)),\n",
    "            ConvBlock(512, 1024, kernel_size=(3, 3), padding=1),\n",
    "            torch.nn.MaxPool2d(kernel_size=(2, 2), stride=2)\n",
    "        )\n",
    "\n",
    "        # BLOCK 5\n",
    "        self.block_5 = torch.nn.Sequential(\n",
    "            *repeat_block([ConvBlock(1024, 512, kernel_size=(1, 1)), ConvBlock(512, 1024, kernel_size=(3, 3), padding=1)], 2),\n",
    "            ConvBlock(1024, 1024, kernel_size=(3, 3), padding=1),\n",
    "            ConvBlock(1024, 1024, kernel_size=(3, 3), stride=2, padding=1)\n",
    "        )\n",
    "\n",
    "        # BLOCK 6\n",
    "        self.block_6 = torch.nn.Sequential(\n",
    "            ConvBlock(1024, 1024, kernel_size=(3, 3), padding='same'),\n",
    "            ConvBlock(1024, 1024, kernel_size=(3, 3), padding='same')\n",
    "        )\n",
    "\n",
    "        # FINAL DENSE\n",
    "        self.final_dense = torch.nn.Sequential(\n",
    "            torch.nn.Flatten(),\n",
    "            torch.nn.Linear(7 * 7 * 1024, self.final_layer_size),\n",
    "            torch.nn.Dropout(0.0),\n",
    "            torch.nn.LeakyReLU(0.1),\n",
    "            # Last layer dims = (S, S, 30) where (C+B*5) = 30\n",
    "            torch.nn.Linear(self.final_layer_size, self.S * self.S * (self.C + self.B * 5))\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.block_1(x)\n",
    "        print(x.shape)\n",
    "        x = self.block_2(x)\n",
    "        print(x.shape)\n",
    "        x = self.block_3(x)\n",
    "        print(x.shape)\n",
    "        x = self.block_4(x)\n",
    "        print(x.shape)\n",
    "        x = self.block_5(x)\n",
    "        print(x.shape)\n",
    "        x = self.block_6(x)\n",
    "        print(x.shape)\n",
    "        x = self.final_dense(x)\n",
    "        print(x.shape)\n",
    "        x = x.view(self.S, self.S, 30)\n",
    "        return x\n",
    "\n",
    "\n",
    "class ConvBlock(torch.nn.Module):\n",
    "\n",
    "    def __init__(self, in_channels: int, out_channels: int, **kwargs):\n",
    "        super().__init__()\n",
    "        self.conv = torch.nn.Conv2d(in_channels, out_channels, **kwargs)\n",
    "        self.bnorm = torch.nn.BatchNorm2d(num_features=out_channels)\n",
    "        self.lrelu = torch.nn.LeakyReLU(negative_slope=0.1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.conv(x)\n",
    "        x = self.bnorm(x)\n",
    "        x = self.lrelu(x)\n",
    "        return x"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 169,
   "outputs": [],
   "source": [
    "model = YoloV1()"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 171,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 64, 112, 112])\n",
      "torch.Size([1, 192, 56, 56])\n",
      "torch.Size([1, 512, 28, 28])\n",
      "torch.Size([1, 1024, 14, 14])\n",
      "torch.Size([1, 1024, 7, 7])\n",
      "torch.Size([1, 1024, 7, 7])\n",
      "torch.Size([1, 1470])\n"
     ]
    },
    {
     "data": {
      "text/plain": "torch.Size([7, 7, 30])"
     },
     "execution_count": 171,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model(xin).shape"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 145,
   "outputs": [],
   "source": [
    "test_b5 = torch.nn.Sequential(\n",
    "            *repeat_block([ConvBlock(1024, 512, kernel_size=(1, 1)), ConvBlock(512, 1024, kernel_size=(3, 3), padding=1)], 2),\n",
    "            ConvBlock(1024, 1024, kernel_size=(3, 3), padding=1),\n",
    "            ConvBlock(1024, 1024, kernel_size=(3, 3), stride=2, padding=1)\n",
    "        )"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 146,
   "outputs": [
    {
     "data": {
      "text/plain": "torch.Size([1, 1024, 7, 7])"
     },
     "execution_count": 146,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_b5(torch.randn((1, 1024, 14, 14))).shape"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}