vgg8s_train.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fully Convolutional Network - Semantic Segmentation\n",
    "\n",
    "![image.png](imgs/2.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![image.png](imgs/3.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import os.path as osp\n",
    "import pytz\n",
    "import torch\n",
    "\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "\n",
    "configurations = {\n",
    "    # same configuration as original work\n",
    "    # https://github.com/shelhamer/fcn.berkeleyvision.org\n",
    "    1: dict(\n",
    "        max_iteration=100000,\n",
    "        lr=1.0e-10,\n",
    "        momentum=0.99,\n",
    "        weight_decay=0.0005,\n",
    "        interval_validate=4000,\n",
    "    )\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from types import SimpleNamespace\n",
    "opts = SimpleNamespace()\n",
    "opts.cfg = configurations[1]\n",
    "opts.resume = ''\n",
    "print(opts.cfg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils import get_log_dir\n",
    "opts.out = get_log_dir('vgg8s', 1, opts.cfg)\n",
    "print(opts.out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gpu = 1\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu)\n",
    "cuda = torch.cuda.is_available()\n",
    "print('Cuda: {}'.format(cuda))\n",
    "opts.cuda = 'cuda' if cuda else 'cpu'\n",
    "opts.mode = 'train'\n",
    "opts.backbone = 'vgg'\n",
    "opts.fcn = '8s'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## PascalVOC Dataset - Downloaded on _`root`_ variable"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "root = './data/Pascal_VOC'\n",
    "print(root)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from data_loader import Pascal_Data\n",
    "kwargs = {'num_workers': 4} if cuda else {}\n",
    "train_loader = torch.utils.data.DataLoader(\n",
    "        Pascal_Data(root, image_set='train', backbone='vgg'),\n",
    "        batch_size=1, shuffle=True, **kwargs)\n",
    "val_loader = torch.utils.data.DataLoader(\n",
    "        Pascal_Data(root, image_set='val', backbone='vgg'),\n",
    "        batch_size=1, shuffle=False, **kwargs)\n",
    "data_loader = [train_loader, val_loader]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "for data, target in train_loader: break\n",
    "print(data.shape)\n",
    "print(target.shape)\n",
    "data.min()\n",
    "data_show, label_show = train_loader.dataset.untransform(data[0].cpu().clone(), target[0].cpu().clone())\n",
    "\n",
    "plt.imshow(data_show)\n",
    "plt.show()\n",
    "\n",
    "def imshow_label(label_show):\n",
    "    import matplotlib\n",
    "    import numpy as np\n",
    "    cmap = plt.cm.jet\n",
    "    # extract all colors from the .jet map\n",
    "    cmaplist = [cmap(i) for i in range(cmap.N)]\n",
    "    cmaplist[0] = (0.0,0.0,0.0,1.0)\n",
    "    cmap = cmap.from_list('Custom cmap', cmaplist, cmap.N)\n",
    "    # define the bins and normalize\n",
    "    bounds = np.arange(0,len(train_loader.dataset.class_names))\n",
    "    norm = matplotlib.colors.BoundaryNorm(bounds, cmap.N)\n",
    "    plt.imshow(label_show, cmap=cmap, norm=norm)\n",
    "    cbar = plt.colorbar(ticks=bounds)\n",
    "    cbar.ax.set_yticklabels(train_loader.dataset.class_names)\n",
    "    plt.show()    \n",
    "    \n",
    "imshow_label(label_show)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## FCN - Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch.nn as nn\n",
    "\n",
    "class FCN8s(nn.Module):\n",
    "\n",
    "    def __init__(self, n_class=21):\n",
    "        super(FCN8s, self).__init__()\n",
    "        # conv1\n",
    "        self.conv1_1 = nn.Conv2d(3, 64, 3, padding=100)\n",
    "        self.relu1_1 = nn.ReLU(inplace=True)\n",
    "        self.conv1_2 = nn.Conv2d(64, 64, 3, padding=1)\n",
    "        self.relu1_2 = nn.ReLU(inplace=True)\n",
    "        self.pool1 = nn.MaxPool2d(2, stride=2, ceil_mode=True)  # 1/2\n",
    "\n",
    "        # conv2\n",
    "        self.conv2_1 = nn.Conv2d(64, 128, 3, padding=1)\n",
    "        self.relu2_1 = nn.ReLU(inplace=True)\n",
    "        self.conv2_2 = nn.Conv2d(128, 128, 3, padding=1)\n",
    "        self.relu2_2 = nn.ReLU(inplace=True)\n",
    "        self.pool2 = nn.MaxPool2d(2, stride=2, ceil_mode=True)  # 1/4\n",
    "\n",
    "        # conv3\n",
    "        self.conv3_1 = nn.Conv2d(128, 256, 3, padding=1)\n",
    "        self.relu3_1 = nn.ReLU(inplace=True)\n",
    "        self.conv3_2 = nn.Conv2d(256, 256, 3, padding=1)\n",
    "        self.relu3_2 = nn.ReLU(inplace=True)\n",
    "        self.conv3_3 = nn.Conv2d(256, 256, 3, padding=1)\n",
    "        self.relu3_3 = nn.ReLU(inplace=True)\n",
    "        self.pool3 = nn.MaxPool2d(2, stride=2, ceil_mode=True)  # 1/8\n",
    "\n",
    "        # conv4\n",
    "        self.conv4_1 = nn.Conv2d(256, 512, 3, padding=1)\n",
    "        self.relu4_1 = nn.ReLU(inplace=True)\n",
    "        self.conv4_2 = nn.Conv2d(512, 512, 3, padding=1)\n",
    "        self.relu4_2 = nn.ReLU(inplace=True)\n",
    "        self.conv4_3 = nn.Conv2d(512, 512, 3, padding=1)\n",
    "        self.relu4_3 = nn.ReLU(inplace=True)\n",
    "        self.pool4 = nn.MaxPool2d(2, stride=2, ceil_mode=True)  # 1/16\n",
    "\n",
    "        # conv5\n",
    "        self.conv5_1 = nn.Conv2d(512, 512, 3, padding=1)\n",
    "        self.relu5_1 = nn.ReLU(inplace=True)\n",
    "        self.conv5_2 = nn.Conv2d(512, 512, 3, padding=1)\n",
    "        self.relu5_2 = nn.ReLU(inplace=True)\n",
    "        self.conv5_3 = nn.Conv2d(512, 512, 3, padding=1)\n",
    "        self.relu5_3 = nn.ReLU(inplace=True)\n",
    "        self.pool5 = nn.MaxPool2d(2, stride=2, ceil_mode=True)  # 1/32\n",
    "\n",
    "        # fc6\n",
    "        self.fc6 = nn.Conv2d(512, 4096, 7)\n",
    "        self.relu6 = nn.ReLU(inplace=True)\n",
    "        self.drop6 = nn.Dropout2d()\n",
    "\n",
    "        # fc7\n",
    "        self.fc7 = nn.Conv2d(4096, 4096, 1)\n",
    "        self.relu7 = nn.ReLU(inplace=True)\n",
    "        self.drop7 = nn.Dropout2d()\n",
    "\n",
    "        self.score_fr = nn.Conv2d(4096, n_class, 1)\n",
    "        self.score_pool3 = nn.Conv2d(256, n_class, 1)\n",
    "        self.score_pool4 = nn.Conv2d(512, n_class, 1)\n",
    "\n",
    "        self.upscore2 = nn.ConvTranspose2d(n_class,\n",
    "                                           n_class,\n",
    "                                           4,\n",
    "                                           stride=2,\n",
    "                                           bias=False)\n",
    "        self.upscore_pool4 = nn.ConvTranspose2d(n_class,\n",
    "                                                n_class,\n",
    "                                                4,\n",
    "                                                stride=2,\n",
    "                                                bias=False)        \n",
    "        self.upscore8 = nn.ConvTranspose2d(n_class,\n",
    "                                           n_class,\n",
    "                                           16,\n",
    "                                           stride=8,\n",
    "                                           bias=False)\n",
    "\n",
    "        self._initialize_weights()\n",
    "\n",
    "    def _initialize_weights(self):\n",
    "        from models.vgg.helpers import get_upsampling_weight\n",
    "        for m in self.modules():\n",
    "            if isinstance(m, nn.Conv2d):\n",
    "                m.weight.data.zero_()\n",
    "                if m.bias is not None:\n",
    "                    m.bias.data.zero_()\n",
    "            if isinstance(m, nn.ConvTranspose2d):\n",
    "                assert m.kernel_size[0] == m.kernel_size[1]\n",
    "                initial_weight = get_upsampling_weight(m.in_channels,\n",
    "                                                       m.out_channels,\n",
    "                                                       m.kernel_size[0])\n",
    "                m.weight.data.copy_(initial_weight)\n",
    "\n",
    "    def forward(self, x, debug=False):\n",
    "        h = x\n",
    "        if debug:\n",
    "            print(h.data.shape)\n",
    "        h = self.relu1_1(self.conv1_1(h))\n",
    "        if debug:\n",
    "            print(h.data.shape)\n",
    "        h = self.relu1_2(self.conv1_2(h))\n",
    "        if debug:\n",
    "            print(h.data.shape)\n",
    "        h = self.pool1(h)\n",
    "        if debug:\n",
    "            print(h.data.shape)\n",
    "\n",
    "        h = self.relu2_1(self.conv2_1(h))\n",
    "        if debug:\n",
    "            print(h.data.shape)\n",
    "        h = self.relu2_2(self.conv2_2(h))\n",
    "        if debug:\n",
    "            print(h.data.shape)\n",
    "        h = self.pool2(h)\n",
    "        if debug:\n",
    "            print(h.data.shape)\n",
    "\n",
    "        h = self.relu3_1(self.conv3_1(h))\n",
    "        if debug:\n",
    "            print(h.data.shape)\n",
    "        h = self.relu3_2(self.conv3_2(h))\n",
    "        if debug:\n",
    "            print(h.data.shape)\n",
    "        h = self.relu3_3(self.conv3_3(h))\n",
    "        if debug:\n",
    "            print(h.data.shape)\n",
    "        h = self.pool3(h)\n",
    "        if debug:\n",
    "            print('pool3: {}'.format(h.data.shape))\n",
    "        pool3 = h  # 1/8\n",
    "\n",
    "        h = self.relu4_1(self.conv4_1(h))\n",
    "        if debug:\n",
    "            print(h.data.shape)\n",
    "        h = self.relu4_2(self.conv4_2(h))\n",
    "        if debug:\n",
    "            print(h.data.shape)\n",
    "        h = self.relu4_3(self.conv4_3(h))\n",
    "        if debug:\n",
    "            print(h.data.shape)\n",
    "        h = self.pool4(h)\n",
    "        if debug:\n",
    "            print('pool4: {}'.format(h.data.shape))\n",
    "        pool4 = h  # 1/16 #<------------------------------------\n",
    "\n",
    "        h = self.relu5_1(self.conv5_1(h))\n",
    "        if debug:\n",
    "            print(h.data.shape)\n",
    "        h = self.relu5_2(self.conv5_2(h))\n",
    "        if debug:\n",
    "            print(h.data.shape)\n",
    "        h = self.relu5_3(self.conv5_3(h))\n",
    "        if debug:\n",
    "            print(h.data.shape)\n",
    "        h = self.pool5(h)\n",
    "        if debug:\n",
    "            print(h.data.shape)\n",
    "\n",
    "        h = self.relu6(self.fc6(h))\n",
    "        if debug:\n",
    "            print(h.data.shape)\n",
    "        h = self.drop6(h)\n",
    "        if debug:\n",
    "            print(h.data.shape)\n",
    "\n",
    "        h = self.relu7(self.fc7(h))\n",
    "        if debug:\n",
    "            print(h.data.shape)\n",
    "        h = self.drop7(h)\n",
    "        if debug:\n",
    "            print(h.data.shape)\n",
    "\n",
    "        h = self.score_fr(h)\n",
    "        if debug:\n",
    "            print(h.data.shape)\n",
    "        h = self.upscore2(h)\n",
    "        if debug:\n",
    "            print('upscore2: {}'.format(h.data.shape))\n",
    "        upscore2 = h  # 1/16\n",
    "\n",
    "        h = self.score_pool4(pool4)\n",
    "        if debug:\n",
    "            print('score_pool4: {}'.format(h.data.shape))\n",
    "        h = h[:, :, 5:5 + upscore2.size()[2], 5:5 + upscore2.size()[3]]\n",
    "        if debug:\n",
    "            print('score_pool4c: {}'.format(h.data.shape))\n",
    "        score_pool4c = h  # 1/16\n",
    "\n",
    "        h = upscore2 + score_pool4c\n",
    "        if debug:\n",
    "            print('upscore2+score_pool4c: {}'.format(h.data.shape))\n",
    "        h = self.upscore_pool4(h)\n",
    "        if debug:\n",
    "            print('upscore_pool4: {}'.format(h.data.shape))        \n",
    "        upscore_pool4 = h  # 1/8\n",
    "\n",
    "        h = self.score_pool3(pool3)\n",
    "        if debug:\n",
    "            print('score_pool3: {}'.format(h.data.shape))\n",
    "        h = h[:, :, 9:9 + upscore_pool4.size()[2], 9:9 +\n",
    "              upscore_pool4.size()[3]]\n",
    "        if debug:\n",
    "            print('score_pool3c: {}'.format(h.data.shape))              \n",
    "        score_pool3c = h  # 1/8\n",
    "\n",
    "        h = upscore_pool4 + score_pool3c  # 1/8\n",
    "        if debug:\n",
    "            print('upscore_pool4+score_pool3c: {}'.format(h.data.shape))\n",
    "\n",
    "        h = self.upscore8(h)\n",
    "        if debug:\n",
    "            print('upscore8: {}'.format(h.data.shape))\n",
    "        h = h[:, :, 31:31 + x.size()[2], 31:31 + x.size()[3]].contiguous()\n",
    "        if debug:\n",
    "            print('upscore8 rearranged: {}'.format(h.data.shape))\n",
    "\n",
    "        return h\n",
    "\n",
    "    def copy_params_from_fcn16s(self, fcn16s):\n",
    "        for name, l1 in fcn16s.named_children():\n",
    "            try:\n",
    "                l2 = getattr(self, name)\n",
    "                l2.weight  # skip ReLU / Dropout\n",
    "            except Exception:\n",
    "                continue\n",
    "            assert l1.weight.size() == l2.weight.size()\n",
    "            l2.weight.data.copy_(l1.weight.data)\n",
    "            if l1.bias is not None:\n",
    "                assert l1.bias.size() == l2.bias.size()\n",
    "                l2.bias.data.copy_(l1.bias.data)             "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# https://github.com/shelhamer/fcn.berkeleyvision.org/blob/master/surgery.py\n",
    "def get_upsampling_weight(in_channels, out_channels, kernel_size):\n",
    "    \"\"\"Make a 2D bilinear kernel suitable for upsampling\"\"\"\n",
    "    factor = (kernel_size + 1) // 2\n",
    "    if kernel_size % 2 == 1:\n",
    "        center = factor - 1\n",
    "    else:\n",
    "        center = factor - 0.5\n",
    "    og = np.ogrid[:kernel_size, :kernel_size]\n",
    "    filt = (1 - abs(og[0] - center) / factor) * \\\n",
    "           (1 - abs(og[1] - center) / factor)\n",
    "    weight = np.zeros((in_channels, out_channels, kernel_size, kernel_size),\n",
    "                      dtype=np.float64)\n",
    "    weight[range(in_channels), range(out_channels), :, :] = filt\n",
    "    return torch.from_numpy(weight).float()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## From fcn16s weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = FCN8s(n_class=21)\n",
    "model.to(opts.cuda)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "iter_loader=iter(train_loader)\n",
    "data, target = next(iter_loader)\n",
    "data = data.to(opts.cuda)\n",
    "with torch.no_grad():\n",
    "    output = model(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('input: ', data.shape)\n",
    "print('output: ', output.data.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data, target = next(iter_loader)\n",
    "data = data.to(opts.cuda)\n",
    "with torch.no_grad():\n",
    "    output = model(data, debug=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data, target = next(iter_loader)\n",
    "data = data.to(opts.cuda)\n",
    "with torch.no_grad():\n",
    "    output = model(data, debug=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if opts.resume:\n",
    "    print('Loading checkpoint from: '+resume)\n",
    "    checkpoint = torch.load(resume)\n",
    "    model.load_state_dict(checkpoint['model_state_dict'])\n",
    "else:\n",
    "    from models.vgg.fcn16s import FCN as FCN16\n",
    "    fcn16s = FCN16()\n",
    "    fcn16s_weights = FCN16.download()  # Original FCN16 pretrained model\n",
    "    fcn16s.load_state_dict(torch.load(fcn16s_weights))\n",
    "    model.copy_params_from_fcn16s(fcn16s)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "from trainer import Trainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = Trainer(data_loader, opts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(opts.cfg.get('interval_validate', len(train_loader))) #Validate every 4000 iterations\n",
    "print(opts.out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start_epoch = 0\n",
    "start_iteration = 0\n",
    "if opts.resume:\n",
    "    start_epoch = checkpoint['epoch']\n",
    "    start_iteration = checkpoint['iteration']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.epoch = start_epoch\n",
    "trainer.iteration = start_iteration\n",
    "trainer.Train()"
   ]
  }
 ],
 "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.7.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}