demo_OpenSetSegmentation_training.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "OpenGAN: Open-Set Recognition via Open Data Generation\n",
    "================\n",
    "**Supplemental Material for ICCV2021 Submission**\n",
    "\n",
    "\n",
    "In this notebook is for demonstrating open-set semantic segmentation, especially for training in this task."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "import packages\n",
    "------------------\n",
    "\n",
    "Some packages are installed automatically through Anaconda. PyTorch should be also installed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3.7.4 (default, Aug 13 2019, 20:35:49) \n",
      "[GCC 7.3.0]\n",
      "1.4.0+cu92\n"
     ]
    }
   ],
   "source": [
    "from __future__ import print_function, division\n",
    "import os, random, time, copy, scipy, pickle, sys, math, json, pickle\n",
    "\n",
    "import argparse, pprint, shutil, logging, time, timeit\n",
    "from pathlib import Path\n",
    "\n",
    "from skimage import io, transform\n",
    "import numpy as np\n",
    "import os.path as path\n",
    "import scipy.io as sio\n",
    "from scipy import misc\n",
    "from scipy import ndimage, signal\n",
    "import matplotlib.pyplot as plt\n",
    "# import PIL.Image\n",
    "from PIL import Image\n",
    "from io import BytesIO\n",
    "from skimage import data, img_as_float\n",
    "from skimage.measure import compare_ssim as ssim\n",
    "from skimage.measure import compare_psnr as psnr\n",
    "\n",
    "import torch, torchvision\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.optim import lr_scheduler \n",
    "import torch.nn.functional as F\n",
    "from torch.autograd import Variable\n",
    "from torchvision import datasets, models, transforms\n",
    "import torchvision.utils as vutils\n",
    "from collections import namedtuple\n",
    "\n",
    "from config_HRNet import models\n",
    "from config_HRNet import seg_hrnet\n",
    "from config_HRNet import config\n",
    "from config_HRNet import update_config\n",
    "from config_HRNet.modelsummary  import *\n",
    "from config_HRNet.utils import *\n",
    "\n",
    "\n",
    "from utils.dataset_tinyimagenet import *\n",
    "from utils.dataset_cityscapes import *\n",
    "from utils.eval_funcs import *\n",
    "\n",
    "\n",
    "import warnings # ignore warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "print(sys.version)\n",
    "print(torch.__version__)\n",
    "\n",
    "# %load_ext autoreload\n",
    "# %autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Setup config parameters\n",
    " -----------------\n",
    " \n",
    " There are several things to setup, like which GPU to use, where to read images and save files, etc. Please read and understand this. By default, you should be able to run this script smoothly by changing nothing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "./exp/demo_step030_OpenGAN_num1000_w0.20\n"
     ]
    }
   ],
   "source": [
    "# set the random seed\n",
    "torch.manual_seed(0)\n",
    "\n",
    "\n",
    "################## set attributes for this project/experiment ##################\n",
    "# config result folder\n",
    "exp_dir = './exp' # experiment directory, used for reading the init model\n",
    "\n",
    "num_open_training_images = 1000\n",
    "weight_adversarialLoss = 0.2\n",
    "project_name = 'demo_step030_OpenGAN_num{}_w{:.2f}'.format(num_open_training_images, weight_adversarialLoss)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "device ='cpu'\n",
    "if torch.cuda.is_available(): \n",
    "    device='cuda:3'\n",
    "        \n",
    "\n",
    "\n",
    "ganBatchSize = 640\n",
    "batch_size = 1\n",
    "newsize = (-1,-1)\n",
    "\n",
    "total_epoch_num = 50 # total number of epoch in training\n",
    "insertConv = False    \n",
    "embDimension = 64\n",
    "#isPretrained = False\n",
    "#encoder_num_layers = 18\n",
    "\n",
    "\n",
    "# Number of channels in the training images. For color images this is 3\n",
    "nc = 720\n",
    "# Size of z latent vector (i.e. size of generator input)\n",
    "nz = 64\n",
    "# Size of feature maps in generator\n",
    "ngf = 64\n",
    "# Size of feature maps in discriminator\n",
    "ndf = 64\n",
    "# Beta1 hyperparam for Adam optimizers\n",
    "beta1 = 0.5\n",
    "# Number of GPUs available. Use 0 for CPU mode.\n",
    "ngpu = 1\n",
    "\n",
    "\n",
    "\n",
    "save_dir = os.path.join(exp_dir, project_name)\n",
    "if not os.path.exists(exp_dir): os.makedirs(exp_dir)\n",
    "\n",
    "lr = 0.0001 # base learning rate\n",
    "\n",
    "num_epochs = total_epoch_num\n",
    "torch.cuda.device_count()\n",
    "torch.cuda.empty_cache()\n",
    "\n",
    "save_dir = os.path.join(exp_dir, project_name)\n",
    "print(save_dir)    \n",
    "if not os.path.exists(save_dir): os.makedirs(save_dir)\n",
    "\n",
    "log_filename = os.path.join(save_dir, 'train.log')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define model architecture\n",
    "---------\n",
    "\n",
    "Here is the definition of the model architecture. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CityscapesOpenPixelFeat4(Dataset):\n",
    "    def __init__(self, set_name='train',\n",
    "                 numImgs=500,\n",
    "                 path_to_data='/scratch/dataset/Cityscapes_feat4'):        \n",
    "        \n",
    "        self.imgList = []\n",
    "        self.current_set_len = numImgs # 2975\n",
    "        if set_name=='test':  \n",
    "            set_name = 'val'\n",
    "            self.current_set_len = 500\n",
    "        \n",
    "        self.set_name = set_name\n",
    "        self.path_to_data = path_to_data\n",
    "        for i in range(self.current_set_len):\n",
    "            self.imgList += ['{}_openpixel.pkl'.format(i)]        \n",
    "        \n",
    "    def __len__(self):        \n",
    "        return self.current_set_len\n",
    "    \n",
    "    def __getitem__(self, idx):        \n",
    "        filename = path.join(self.path_to_data, self.set_name, self.imgList[idx])\n",
    "        with open(filename, \"rb\") as fn:\n",
    "            openPixFeat = pickle.load(fn)\n",
    "        openPixFeat = openPixFeat['feat4open_percls']\n",
    "        openPixFeat = torch.cat(openPixFeat, 0).detach()\n",
    "        #print(openPixFeat.shape)\n",
    "        return openPixFeat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "parser = argparse.ArgumentParser(description='Train segmentation network') \n",
    "parser.add_argument('--cfg',\n",
    "                    help='experiment configure file name',\n",
    "                    default='./config_HRNet/seg_hrnet_w48_train_512x1024_sgd_lr1e-2_wd5e-4_bs_12_epoch484.yaml',\n",
    "                    type=str)\n",
    "parser.add_argument('opts',\n",
    "                    help=\"Modify config options using the command-line\",\n",
    "                    default=None,\n",
    "                    nargs=argparse.REMAINDER)\n",
    "\n",
    "\n",
    "args = parser.parse_args(r'--cfg  ./config_HRNet/seg_hrnet_w48_train_512x1024_sgd_lr1e-2_wd5e-4_bs_12_epoch484.yaml ')\n",
    "args.opts = []\n",
    "update_config(config, args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = eval(config.MODEL.NAME + '.get_seg_model_myModel')(config)\n",
    "model_dict = model.state_dict()\n",
    "\n",
    "\n",
    "model_state_file = '../openset/models/hrnet_w48_cityscapes_cls19_1024x2048_ohem_trainset.pth'\n",
    "pretrained_dict = torch.load(model_state_file, map_location=lambda storage, loc: storage)\n",
    "\n",
    "\n",
    "suppl_dict = {}\n",
    "suppl_dict['last_1_conv.weight'] = pretrained_dict['model.last_layer.0.weight'].clone()\n",
    "suppl_dict['last_1_conv.bias'] = pretrained_dict['model.last_layer.0.bias'].clone()\n",
    "\n",
    "suppl_dict['last_2_BN.running_mean'] = pretrained_dict['model.last_layer.1.running_mean'].clone()\n",
    "suppl_dict['last_2_BN.running_var'] = pretrained_dict['model.last_layer.1.running_var'].clone()\n",
    "# suppl_dict['last_2_BN.num_batches_tracked'] = pretrained_dict['model.last_layer.1.num_batches_tracked']\n",
    "suppl_dict['last_2_BN.weight'] = pretrained_dict['model.last_layer.1.weight'].clone()\n",
    "suppl_dict['last_2_BN.bias'] = pretrained_dict['model.last_layer.1.bias'].clone()\n",
    "\n",
    "suppl_dict['last_4_conv.weight'] = pretrained_dict['model.last_layer.3.weight'].clone()\n",
    "suppl_dict['last_4_conv.bias'] = pretrained_dict['model.last_layer.3.bias'].clone()\n",
    "\n",
    "\n",
    "pretrained_dict = {k[6:]: v for k, v in pretrained_dict.items()\n",
    "                   if k[6:] in model_dict.keys()}\n",
    "\n",
    "\n",
    "model_dict.update(pretrained_dict)\n",
    "model_dict.update(suppl_dict)\n",
    "model.load_state_dict(model_dict)\n",
    "\n",
    "\n",
    "model.eval();\n",
    "model.to(device);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def weights_init(m):\n",
    "    classname = m.__class__.__name__\n",
    "    if classname.find('Conv') != -1:\n",
    "        nn.init.normal_(m.weight.data, 0.0, 0.02)\n",
    "    elif classname.find('BatchNorm') != -1:\n",
    "        nn.init.normal_(m.weight.data, 1.0, 0.02)\n",
    "        nn.init.constant_(m.bias.data, 0)     \n",
    "        \n",
    "\n",
    "class Generator(nn.Module):\n",
    "    def __init__(self, ngpu=1, nz=100, ngf=64, nc=512):\n",
    "        super(Generator, self).__init__()\n",
    "        self.ngpu = ngpu\n",
    "        self.nz = nz\n",
    "        self.ngf = ngf\n",
    "        self.nc = nc\n",
    "        \n",
    "        self.main = nn.Sequential(\n",
    "            # input is Z, going into a convolution\n",
    "            nn.Conv2d( self.nz, self.ngf * 8, 1, 1, 0, bias=True),\n",
    "            nn.BatchNorm2d(self.ngf * 8),\n",
    "            nn.ReLU(True),\n",
    "            # state size. (self.ngf*8) x 4 x 4\n",
    "            nn.Conv2d(self.ngf * 8, self.ngf * 4, 1, 1, 0, bias=True),\n",
    "            nn.BatchNorm2d(self.ngf * 4),\n",
    "            nn.ReLU(True),\n",
    "            # state size. (self.ngf*4) x 8 x 8\n",
    "            nn.Conv2d( self.ngf * 4, self.ngf * 2, 1, 1, 0, bias=True),\n",
    "            nn.BatchNorm2d(self.ngf * 2),\n",
    "            nn.ReLU(True),\n",
    "            # state size. (self.ngf*2) x 16 x 16\n",
    "            nn.Conv2d( self.ngf * 2, self.ngf*4, 1, 1, 0, bias=True),\n",
    "            nn.BatchNorm2d(self.ngf*4),\n",
    "            nn.ReLU(True),\n",
    "            # state size. (self.ngf) x 32 x 32\n",
    "            nn.Conv2d( self.ngf*4, self.nc, 1, 1, 0, bias=True),\n",
    "            #nn.Tanh()\n",
    "            # state size. (self.nc) x 64 x 64\n",
    "        )\n",
    "\n",
    "    def forward(self, input):\n",
    "        return self.main(input)\n",
    "\n",
    "    \n",
    "class Discriminator(nn.Module):\n",
    "    def __init__(self, ngpu=1, nc=512, ndf=64):\n",
    "        super(Discriminator, self).__init__()\n",
    "        self.ngpu = ngpu\n",
    "        self.nc = nc\n",
    "        self.ndf = ndf\n",
    "        self.main = nn.Sequential(\n",
    "            nn.Conv2d(self.nc, self.ndf*8, 1, 1, 0, bias=True),\n",
    "            nn.LeakyReLU(0.2, inplace=True),\n",
    "            nn.Conv2d(self.ndf*8, self.ndf*4, 1, 1, 0, bias=True),\n",
    "            nn.BatchNorm2d(self.ndf*4),\n",
    "            nn.LeakyReLU(0.2, inplace=True),\n",
    "            nn.Conv2d(self.ndf*4, self.ndf*2, 1, 1, 0, bias=True),\n",
    "            nn.BatchNorm2d(self.ndf*2),\n",
    "            nn.LeakyReLU(0.2, inplace=True),\n",
    "            nn.Conv2d(self.ndf*2, self.ndf, 1, 1, 0, bias=True),\n",
    "            nn.BatchNorm2d(self.ndf),\n",
    "            nn.LeakyReLU(0.2, inplace=True),\n",
    "            nn.Conv2d(self.ndf, 1, 1, 1, 0, bias=True),\n",
    "            nn.Sigmoid()\n",
    "        )\n",
    "\n",
    "    def forward(self, input):\n",
    "        return self.main(input)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cuda:3\n"
     ]
    }
   ],
   "source": [
    "netG = Generator(ngpu=ngpu, nz=nz, ngf=ngf, nc=nc).to(device)\n",
    "netD = Discriminator(ngpu=ngpu, nc=nc, ndf=ndf).to(device)\n",
    "\n",
    "\n",
    "# Handle multi-gpu if desired\n",
    "if ('cuda' in device) and (ngpu > 1): \n",
    "    netD = nn.DataParallel(netD, list(range(ngpu)))\n",
    "\n",
    "# Apply the weights_init function to randomly initialize all weights\n",
    "#  to mean=0, stdev=0.2.\n",
    "netD.apply(weights_init)\n",
    "\n",
    "\n",
    "if ('cuda' in device) and (ngpu > 1):\n",
    "    netG = nn.DataParallel(netG, list(range(ngpu)))\n",
    "netG.apply(weights_init)\n",
    "\n",
    "print(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([5, 64, 1, 1]) torch.Size([5, 720, 1, 1]) torch.Size([5, 1, 1, 1])\n"
     ]
    }
   ],
   "source": [
    "noise = torch.randn(batch_size*5, nz, 1, 1, device=device)\n",
    "# Generate fake image batch with G\n",
    "fake = netG(noise)\n",
    "predLabel = netD(fake)\n",
    "\n",
    "print(noise.shape, fake.shape, predLabel.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "setup dataset\n",
    "-----------"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2975 500\n"
     ]
    }
   ],
   "source": [
    "# torchvision.transforms.Normalize(mean, std, inplace=False)\n",
    "imgTransformList = transforms.Compose([\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),\n",
    "])\n",
    "\n",
    "targetTransformList = transforms.Compose([\n",
    "    transforms.ToTensor(),    \n",
    "])\n",
    "\n",
    "cls_datasets = {set_name: Cityscapes(root='/scratch/dataset/Cityscapes',\n",
    "                                     newsize=newsize,\n",
    "                                     split=set_name,\n",
    "                                     mode='fine',\n",
    "                                     target_type='semantic',\n",
    "                                     transform=imgTransformList,\n",
    "                                     target_transform=targetTransformList,\n",
    "                                     transforms=None)\n",
    "                for set_name in ['train', 'val']} # 'train', \n",
    "\n",
    "dataloaders = {set_name: DataLoader(cls_datasets[set_name],\n",
    "                                    batch_size=batch_size,\n",
    "                                    shuffle=set_name=='train', \n",
    "                                    num_workers=4) # num_work can be set to batch_size\n",
    "               for set_name in ['train', 'val']} # 'train',\n",
    "\n",
    "\n",
    "print(len(cls_datasets['train']), len(cls_datasets['val']))\n",
    "classDictionary = cls_datasets['val'].classes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 unlabeled\n",
      "1 ego vehicle\n",
      "2 rectification border\n",
      "3 out of roi\n",
      "4 static\n",
      "5 dynamic\n",
      "6 ground\n",
      "9 parking\n",
      "10 rail track\n",
      "14 guard rail\n",
      "15 bridge\n",
      "16 tunnel\n",
      "18 polegroup\n",
      "29 caravan\n",
      "30 trailer\n",
      "34 license plate\n",
      "total# 16\n"
     ]
    }
   ],
   "source": [
    "id2trainID = {}\n",
    "id2color = {}\n",
    "trainID2color = {}\n",
    "id2name = {}\n",
    "opensetIDlist = []\n",
    "for i in range(len(classDictionary)):\n",
    "    id2trainID[i] = classDictionary[i][2]\n",
    "    id2color[i] = classDictionary[i][-1]\n",
    "    trainID2color[classDictionary[i][2]] = classDictionary[i][-1]\n",
    "    id2name[i] = classDictionary[i][0]\n",
    "    if classDictionary[i][-2]:\n",
    "        opensetIDlist += [i]\n",
    "\n",
    "id2trainID_list = []\n",
    "for i in range(len(id2trainID)):\n",
    "    id2trainID_list.append(id2trainID[i])\n",
    "id2trainID_np = np.asarray(id2trainID_list)        \n",
    "        \n",
    "for elm in opensetIDlist:\n",
    "    print(elm, id2name[elm])\n",
    "print('total# {}'.format(len(opensetIDlist)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_sampler = iter(dataloaders['train'])\n",
    "data = next(data_sampler)\n",
    "imageList, labelList = data[0], data[1]\n",
    "\n",
    "imageList = imageList.to(device)\n",
    "labelList = labelList.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([1, 3, 1024, 2048]), torch.Size([1, 1024, 2048]))"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "imageList.shape, labelList.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "setup training\n",
    "-----------"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize BCELoss function\n",
    "criterion = nn.BCELoss()\n",
    "\n",
    "# Create batch of latent vectors that we will use to visualize\n",
    "#  the progression of the generator\n",
    "fixed_noise = torch.randn(64, nz, 1, 1, device=device)\n",
    "\n",
    "# Establish open and close labels\n",
    "close_label = 1\n",
    "open_label = 0\n",
    "\n",
    "# Establish convention for real and fake labels during training\n",
    "real_label = 1\n",
    "fake_label = 0\n",
    "\n",
    "\n",
    "# Setup Adam optimizers for both G and D\n",
    "optimizerD = optim.Adam(netD.parameters(), lr=lr/1.5, betas=(beta1, 0.999))\n",
    "optimizerG = optim.Adam(netG.parameters(), lr=lr, betas=(beta1, 0.999))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "testing a single image\n",
    "-----------"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "labelList = labelList.unsqueeze(1)\n",
    "labelList = F.interpolate(labelList, scale_factor=0.25, mode='nearest')\n",
    "labelList = labelList.squeeze()\n",
    "H, W = labelList.squeeze().shape\n",
    "trainlabelList = id2trainID_np[labelList.cpu().numpy().reshape(-1,).astype(np.int32)]\n",
    "trainlabelList = trainlabelList.reshape((1,H,W))\n",
    "trainlabelList = torch.from_numpy(trainlabelList)\n",
    "\n",
    "\n",
    "\n",
    "upsampleFunc = nn.UpsamplingBilinear2d(scale_factor=4)\n",
    "with torch.no_grad():\n",
    "    imageList = imageList.to(device)\n",
    "    logitsTensor = model(imageList).detach().cpu()\n",
    "    #logitsTensor = upsampleFunc(logitsTensor)\n",
    "    softmaxTensor = F.softmax(logitsTensor, dim=1)\n",
    "    \n",
    "    feat1Tensor = model.feat1.detach()\n",
    "    feat2Tensor = model.feat2.detach()\n",
    "    feat3Tensor = model.feat3.detach()\n",
    "    feat4Tensor = model.feat4.detach()\n",
    "    feat5Tensor = model.feat5.detach()\n",
    "    \n",
    "    torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([1, 720, 256, 512]), torch.Size([1, 256, 512]), 131072)"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "feat4Tensor.shape, trainlabelList.shape, trainlabelList.shape[1]*trainlabelList.shape[2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "validList = trainlabelList.reshape(-1,1)\n",
    "validList = ((validList>=0) & (validList<=18)).nonzero()\n",
    "validList = validList[:,0]\n",
    "validList = validList[torch.randperm(validList.size()[0])]\n",
    "validList = validList[:ganBatchSize]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "label = torch.full((ganBatchSize,), close_label, device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "real_cpu = feat4Tensor.squeeze()\n",
    "real_cpu = real_cpu.reshape(real_cpu.shape[0], -1).permute(1,0)\n",
    "real_cpu = real_cpu[validList,:].unsqueeze(-1).unsqueeze(-1).to(device)\n",
    "\n",
    "output = netD(real_cpu).view(-1)\n",
    "# Calculate loss on all-real batch\n",
    "errD_real = criterion(output, label)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "noise = torch.randn(ganBatchSize, nz, 1, 1, device=device)\n",
    "# Generate fake image batch with G\n",
    "fake = netG(noise)\n",
    "label.fill_(fake_label)\n",
    "# Classify all fake batch with D\n",
    "output = netD(fake.detach()).view(-1)\n",
    "# Calculate D's loss on the all-fake batch\n",
    "errD_fake = criterion(output, label)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([640, 64, 1, 1]), torch.Size([640]), torch.Size([640, 720, 1, 1]))"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "noise.shape, label.shape, fake.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "training GAN\n",
    "-----------"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([640, 720, 1, 1])\n"
     ]
    }
   ],
   "source": [
    "openPix_datasets = CityscapesOpenPixelFeat4(set_name='train', numImgs=num_open_training_images)\n",
    "openPix_dataloader = DataLoader(openPix_datasets, batch_size=1, shuffle=True, num_workers=4)               \n",
    "\n",
    "openPix_sampler = iter(openPix_dataloader)\n",
    "\n",
    "openPixFeat = next(openPix_sampler)\n",
    "openPixFeat = openPixFeat.squeeze(0)\n",
    "\n",
    "openPixIdxList = torch.randperm(openPixFeat.size()[0])\n",
    "openPixIdxList = openPixIdxList[:ganBatchSize]\n",
    "openPixFeat = openPixFeat[openPixIdxList].to(device)\n",
    "\n",
    "print(openPixFeat.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starting Training Loop...\n",
      "[0/50][0/2975]\t\tlossG: 0.6536, lossD: 0.5840\n",
      "[0/50][100/2975]\t\tlossG: 0.6569, lossD: 0.4096\n",
      "[0/50][200/2975]\t\tlossG: 0.6236, lossD: 0.3967\n",
      "[0/50][300/2975]\t\tlossG: 0.5869, lossD: 0.2977\n",
      "[0/50][400/2975]\t\tlossG: 0.5548, lossD: 0.2332\n",
      "[0/50][500/2975]\t\tlossG: 0.5504, lossD: 0.3024\n"
     ]
    }
   ],
   "source": [
    "# Training Loop\n",
    "\n",
    "# Lists to keep track of progress\n",
    "lossList = []\n",
    "G_losses = []\n",
    "D_losses = []\n",
    "\n",
    "fake_BatchSize = int(ganBatchSize/2)\n",
    "open_BatchSize = ganBatchSize\n",
    "\n",
    "\n",
    "\n",
    "tmp_weights = torch.full((ganBatchSize+open_BatchSize+fake_BatchSize,), 1, device=device)\n",
    "tmp_weights[-fake_BatchSize:] *= weight_adversarialLoss\n",
    "criterionD = nn.BCELoss(weight=tmp_weights)\n",
    "\n",
    "\n",
    "\n",
    "print(\"Starting Training Loop...\")\n",
    "# For each epoch\n",
    "openPixImgCount = 0\n",
    "openPix_sampler = iter(openPix_dataloader)\n",
    "for epoch in range(num_epochs):\n",
    "    # For each batch in the dataloader\n",
    "    for i, sample in enumerate(dataloaders['train'], 0):\n",
    "        imageList, labelList = sample\n",
    "        imageList = imageList.to(device)\n",
    "        labelList = labelList.to(device)\n",
    "\n",
    "        labelList = labelList.unsqueeze(1)\n",
    "        labelList = F.interpolate(labelList, scale_factor=0.25, mode='nearest')\n",
    "        labelList = labelList.squeeze()\n",
    "        H, W = labelList.squeeze().shape\n",
    "        trainlabelList = id2trainID_np[labelList.cpu().numpy().reshape(-1,).astype(np.int32)]\n",
    "        trainlabelList = trainlabelList.reshape((1,H,W))\n",
    "        trainlabelList = torch.from_numpy(trainlabelList)\n",
    "        \n",
    "        \n",
    "        #upsampleFunc = nn.UpsamplingBilinear2d(scale_factor=4)\n",
    "        with torch.no_grad():\n",
    "            imageList = imageList.to(device)\n",
    "            logitsTensor = model(imageList).detach().cpu()\n",
    "            featTensor = model.feat4.detach()\n",
    "        \n",
    "        validList = trainlabelList.reshape(-1,1)\n",
    "        validList = ((validList>=0) & (validList<=18)).nonzero()\n",
    "        validList = validList[:,0]\n",
    "        tmp = torch.randperm(validList.size()[0])        \n",
    "        validList = validList[tmp[:ganBatchSize]]\n",
    "                \n",
    "\n",
    "        \n",
    "        label_closeset = torch.full((ganBatchSize,), close_label, device=device)\n",
    "        feat_closeset = featTensor.squeeze()\n",
    "        feat_closeset = feat_closeset.reshape(feat_closeset.shape[0], -1).permute(1,0)\n",
    "        feat_closeset = feat_closeset[validList,:].unsqueeze(-1).unsqueeze(-1)        \n",
    "        label_open = torch.full((open_BatchSize,), open_label, device=device)\n",
    "        \n",
    "        openPixImgCount += 1\n",
    "        feat_openset = next(openPix_sampler)\n",
    "        feat_openset = feat_openset.squeeze(0)\n",
    "        openPixIdxList = torch.randperm(feat_openset.size()[0])\n",
    "        openPixIdxList = openPixIdxList[:open_BatchSize]\n",
    "        feat_openset = feat_openset[openPixIdxList].to(device)\n",
    "\n",
    "        if openPixImgCount==num_open_training_images:\n",
    "            openPixImgCount = 0\n",
    "            openPix_sampler = iter(openPix_dataloader)\n",
    "        \n",
    "        \n",
    "        \n",
    "        # generate fake images        \n",
    "        noise = torch.randn(fake_BatchSize, nz, 1, 1, device=device)\n",
    "        # Generate fake image batch with G\n",
    "        label_fake = torch.full((fake_BatchSize,), fake_label, device=device)\n",
    "        feat_fakeset = netG(noise)    \n",
    "        \n",
    "        ############################\n",
    "        # (1) Update D network: maximize log(D(x)) + log(1 - D(G(z)))\n",
    "        ###########################\n",
    "        # using close&open&fake data to update D\n",
    "        netD.zero_grad()\n",
    "        X = torch.cat((feat_closeset, feat_openset.to(device), feat_fakeset.detach()),0)\n",
    "        label_total = torch.cat((label_closeset, label_open, label_fake),0)\n",
    "                \n",
    "        output = netD(X).view(-1)\n",
    "        lossD = criterionD(output, label_total)\n",
    "        lossD.backward()\n",
    "        optimizerD.step()\n",
    "        errD = lossD.mean().item()                        \n",
    "            \n",
    "            \n",
    "        ############################\n",
    "        # (2) Update G network: maximize log(D(G(z)))\n",
    "        ###########################\n",
    "        netG.zero_grad()\n",
    "        label_fakeclose = torch.full((fake_BatchSize,), close_label, device=device)        \n",
    "        # Since we just updated D, perform another forward pass of all-fake batch through D\n",
    "        output = netD(feat_fakeset).view(-1)\n",
    "        # Calculate G's loss based on this output\n",
    "        lossG = criterion(output, label_fakeclose)\n",
    "        # Calculate gradients for G\n",
    "        lossG.backward()\n",
    "        errG = lossG.mean().item()\n",
    "        # Update G\n",
    "        optimizerG.step()\n",
    "            \n",
    "            \n",
    "        # Save Losses for plotting later\n",
    "        G_losses.append(errG)\n",
    "        D_losses.append(errD)\n",
    "        \n",
    "        \n",
    "        # Output training stats\n",
    "        if i % 100 == 0:\n",
    "            print('[%d/%d][%d/%d]\\t\\tlossG: %.4f, lossD: %.4f'\n",
    "                  % (epoch, num_epochs, i, len(dataloaders['train']), \n",
    "                     errG, errD))\n",
    "            \n",
    "            \n",
    "    cur_model_wts = copy.deepcopy(netD.state_dict())\n",
    "    path_to_save_paramOnly = os.path.join(save_dir, 'epoch-{}.classifier'.format(epoch+1))\n",
    "    torch.save(cur_model_wts, path_to_save_paramOnly)\n",
    "    cur_model_wts = copy.deepcopy(netG.state_dict())\n",
    "    path_to_save_paramOnly = os.path.join(save_dir, 'epoch-{}.GNet'.format(epoch+1))\n",
    "    torch.save(cur_model_wts, path_to_save_paramOnly)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "validating results\n",
    "-----------"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(10,5))\n",
    "plt.title(\"binary cross-entropy loss in training\")\n",
    "plt.plot(Dopen_losses, label=\"Dopen_losses\")\n",
    "plt.plot(Dclose_losses, label=\"Dclose_losses\")\n",
    "plt.plot(Dfake_losses, label=\"Dfake_losses\")\n",
    "plt.plot(G_losses, label=\"G_losses\")\n",
    "plt.xlabel(\"iterations\")\n",
    "plt.ylabel(\"Loss\")\n",
    "plt.legend()\n",
    "# plt.savefig('learningCurves_{}.png'.format(modelFlag), bbox_inches='tight',transparent=True)\n",
    "# plt.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.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}