Test-gating.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "horizontal-recovery",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "\n",
    "from pathlib import Path\n",
    "\n",
    "from data.dataset import Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "according-breakfast",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datetime import datetime\n",
    "\n",
    "# datetime object containing current date and time\n",
    "now = datetime.now()\n",
    " \n",
    "print(\"now =\", now)\n",
    "\n",
    "# dd/mm/YY H:M:S\n",
    "dt_string = now.strftime(\"%d-%m-%Y_%H.%M.%S\")\n",
    "print(\"date and time =\", dt_string)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fatty-copying",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(torch.__version__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "listed-radical",
   "metadata": {},
   "outputs": [],
   "source": [
    "# print(Path.home())\n",
    "base_path = f'{Path.home()}/SageMaker'\n",
    "print(base_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "encouraging-rebound",
   "metadata": {},
   "outputs": [],
   "source": [
    "# train_dataset = Dataset(data_path=config['train_data_path'],\n",
    "#                                 with_subfolder=False,\n",
    "#                                 image_shape=config['image_shape'],\n",
    "#                                 random_crop=config['random_crop'])\n",
    "\n",
    "dataset = Dataset(data_path=f'{base_path}/dataset/flickr8k/Images',\n",
    "                                with_subfolder=False,\n",
    "                                image_shape=[256,256,3],\n",
    "                                random_crop=True)\n",
    "\n",
    "loader = torch.utils.data.DataLoader(dataset, \n",
    "                                     batch_size=4,\n",
    "                                     shuffle=True,\n",
    "                                     num_workers=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "published-arrow",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "# functions to show an image\n",
    "\n",
    "\n",
    "def imshow(img):\n",
    "    img = img / 2 + 0.5     # unnormalize\n",
    "    npimg = img.numpy()\n",
    "    plt.imshow(np.transpose(npimg, (1, 2, 0)))\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "# get some random training images\n",
    "dataiter = iter(loader)\n",
    "images = dataiter.next()\n",
    "\n",
    "# show images\n",
    "imshow(torchvision.utils.make_grid(images))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "amino-freeware",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "helpful-election",
   "metadata": {},
   "source": [
    "## Instantiate the model, and test gating."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ceramic-hudson",
   "metadata": {},
   "outputs": [],
   "source": [
    "from trainer import Trainer\n",
    "from data.dataset import Dataset\n",
    "from utils.tools import get_config, random_bbox, mask_image\n",
    "from model.networks import Generator, LocalDis, GlobalDis\n",
    "from utils.tools import get_model_list, local_patch, spatial_discounting_mask\n",
    "\n",
    "# from utils.logger import get_logger\n",
    "\n",
    "from argparse import ArgumentParser\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "editorial-george",
   "metadata": {},
   "outputs": [],
   "source": [
    "parser = ArgumentParser()\n",
    "parser.add_argument('--config', type=str, default='configs/config.yaml',\n",
    "                    help=\"training configuration\")\n",
    "parser.add_argument('--seed', type=int, help='manual seed')\n",
    "\n",
    "args = parser.parse_args('')\n",
    "config = get_config(args.config)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "altered-groove",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "metallic-arabic",
   "metadata": {},
   "source": [
    "## Grab some data and prepare masks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "optical-wages",
   "metadata": {},
   "outputs": [],
   "source": [
    "# try:\n",
    "#     ground_truth = next(loader)\n",
    "# except StopIteration:\n",
    "#     iterable_train_loader = iter(train_loader)\n",
    "#     ground_truth = next(iterable_train_loader)\n",
    "\n",
    "ground_truth = images\n",
    "\n",
    "# Prepare the inputs\n",
    "bboxes = random_bbox(config, batch_size=ground_truth.size(0))\n",
    "x, mask = mask_image(ground_truth, bboxes, config)\n",
    "if config['cuda']:\n",
    "    x = x.cuda()\n",
    "    mask = mask.cuda()\n",
    "    ground_truth = ground_truth.cuda()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "celtic-practice",
   "metadata": {},
   "source": [
    "## Generator - non-gated"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "absent-filter",
   "metadata": {},
   "outputs": [],
   "source": [
    "from model.networks import Generator\n",
    "# del netG\n",
    "\n",
    "config = get_config(f'{base_path}/generative-inpainting-pytorch/configs/config.yaml')\n",
    "netG = Generator(config['netG'], config['cuda'], config['gpu_ids']).cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "quiet-search",
   "metadata": {},
   "outputs": [],
   "source": [
    "x1, x2, offset_flow = netG(x, mask)\n",
    "print(\"x1: \", x1.shape)\n",
    "print(\"x2: \", x2.shape)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "expired-antarctica",
   "metadata": {},
   "source": [
    "## Inspect individual outputs from layers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "clinical-canada",
   "metadata": {},
   "outputs": [],
   "source": [
    "x1 = netG.coarse_generator(x, mask)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "controlling-torture",
   "metadata": {},
   "outputs": [],
   "source": [
    "ones = torch.ones(x.size(0), 1, x.size(2), x.size(3)).cuda()\n",
    "print(\"ones: \", ones.shape)\n",
    "\n",
    "_in = torch.cat([x, ones, mask], dim=1)\n",
    "print(\"_in: \", _in.shape)\n",
    "_x1 = netG.coarse_generator.conv1(_in)\n",
    "print(\"_x1: \", _x1.shape)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "threatened-convenience",
   "metadata": {},
   "source": [
    "## Gated Generator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "embedded-summer",
   "metadata": {},
   "outputs": [],
   "source": [
    "from model.networks import Generator\n",
    "del netG_gated\n",
    "\n",
    "config_gated = get_config(f'{base_path}/generative-inpainting-pytorch/configs/config-gated.yaml')\n",
    "netG_gated = Generator(config_gated['netG'], config_gated['cuda'], config_gated['gpu_ids']).cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "virtual-safety",
   "metadata": {},
   "outputs": [],
   "source": [
    "# x1, x2, offset_flow = netG_gated(x, mask)\n",
    "print(\"x:, \", x.shape, \" mask: \", mask.shape)\n",
    "x1 = netG_gated.coarse_generator(x, mask)\n",
    "print(\"x1: \", x1.shape)\n",
    "# print(\"x2: \", x2.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "surrounded-surfing",
   "metadata": {},
   "outputs": [],
   "source": [
    "ones = torch.ones(x.size(0), 1, x.size(2), x.size(3)).cuda()\n",
    "print(\"ones: \", ones.shape)\n",
    "\n",
    "_in = torch.cat([x, ones, mask], dim=1)\n",
    "print(\"_in: \", _in.shape)\n",
    "_x1 = netG_gated.coarse_generator.conv1(_in)\n",
    "print(\"_x1: \", _x1.shape)\n",
    "\n",
    "_x2 = netG_gated.coarse_generator.conv2_downsample(_x1)\n",
    "print(\"_x2: \", _x2.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "addressed-saver",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "silver-sister",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "listed-event",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "\n",
    "class Conv2dBlockGated(nn.Module):\n",
    "    def __init__(self, input_dim, output_dim, kernel_size, stride, padding=0,\n",
    "                 conv_padding=0, dilation=1, weight_norm='none', norm='none',\n",
    "                 activation='relu', pad_type='zero', transpose=False):\n",
    "        \n",
    "        super(Conv2dBlockGated, self).__init__()\n",
    "        self.output_dim = output_dim\n",
    "        self.use_bias = True\n",
    "        # initialize padding\n",
    "        if pad_type == 'reflect':\n",
    "            self.pad = nn.ReflectionPad2d(padding)\n",
    "        elif pad_type == 'replicate':\n",
    "            self.pad = nn.ReplicationPad2d(padding)\n",
    "        elif pad_type == 'zero':\n",
    "            self.pad = nn.ZeroPad2d(padding)\n",
    "        elif pad_type == 'none':\n",
    "            self.pad = None\n",
    "        else:\n",
    "            assert 0, \"Unsupported padding type: {}\".format(pad_type)\n",
    "\n",
    "        # initialize normalization\n",
    "        norm_dim = output_dim\n",
    "        if norm == 'bn':\n",
    "            self.norm = nn.BatchNorm2d(norm_dim)\n",
    "        elif norm == 'in':\n",
    "            self.norm = nn.InstanceNorm2d(norm_dim)\n",
    "        elif norm == 'none':\n",
    "            self.norm = None\n",
    "        else:\n",
    "            assert 0, \"Unsupported normalization: {}\".format(norm)\n",
    "\n",
    "        if weight_norm == 'sn':\n",
    "            self.weight_norm = spectral_norm_fn\n",
    "        elif weight_norm == 'wn':\n",
    "            self.weight_norm = weight_norm_fn\n",
    "        elif weight_norm == 'none':\n",
    "            self.weight_norm = None\n",
    "        else:\n",
    "            assert 0, \"Unsupported normalization: {}\".format(weight_norm)\n",
    "\n",
    "        # initialize activation\n",
    "        if activation == 'relu':\n",
    "            self.activation = nn.ReLU(inplace=True)\n",
    "        elif activation == 'elu':\n",
    "            self.activation = nn.ELU(inplace=True)\n",
    "        elif activation == 'lrelu':\n",
    "            self.activation = nn.LeakyReLU(0.2, inplace=True)\n",
    "        elif activation == 'prelu':\n",
    "            self.activation = nn.PReLU()\n",
    "        elif activation == 'selu':\n",
    "            self.activation = nn.SELU(inplace=True)\n",
    "        elif activation == 'tanh':\n",
    "            self.activation = nn.Tanh()\n",
    "        elif activation == 'none':\n",
    "            self.activation = None\n",
    "        else:\n",
    "            assert 0, \"Unsupported activation: {}\".format(activation)\n",
    "\n",
    "        # initialize convolution\n",
    "        if transpose:\n",
    "            self.conv = nn.ConvTranspose2d(input_dim, output_dim,\n",
    "                                           kernel_size, stride,\n",
    "                                           padding=conv_padding,\n",
    "                                           output_padding=conv_padding,\n",
    "                                           dilation=dilation,\n",
    "                                           bias=self.use_bias)\n",
    "\n",
    "        else:\n",
    "            self.conv = nn.Conv2d(input_dim, output_dim, kernel_size, stride,\n",
    "                                  padding=conv_padding, dilation=dilation,\n",
    "                                  bias=self.use_bias)\n",
    "\n",
    "            \n",
    "\n",
    "        if self.weight_norm:\n",
    "            self.conv = self.weight_norm(self.conv)\n",
    "            \n",
    "\n",
    "    def forward(self, x):\n",
    "        if self.pad:\n",
    "            x = self.conv(self.pad(x))\n",
    "        else:\n",
    "            x = self.conv(x)\n",
    "            \n",
    "            \n",
    "        if self.norm:\n",
    "            x = self.norm(x)\n",
    "            \n",
    "            \n",
    "        # If there are more than 3 channels, then we treat the remainder as the mask and optional input\n",
    "        # and \"gate\" that.\n",
    "        #\n",
    "        print(\"x: \", x.shape)\n",
    "        feat, gate = torch.chunk(x, 2, 1)\n",
    "        print(\"feat: \", feat.shape)\n",
    "        print(\"gate: \", gate.shape)\n",
    "        \n",
    "        \n",
    "        # Output is image or no activation.\n",
    "        #\n",
    "        if self.activation is None or self.output_dim == 3:\n",
    "            return x\n",
    "        \n",
    "        \n",
    "        # Otherwise we compute activation of features and gate.\n",
    "        #\n",
    "        feat = self.activation(feat)\n",
    "        gate = torch.sigmoid(gate) # Gate\n",
    "    \n",
    "        return feat * gate\n",
    "            \n",
    "\n",
    "\n",
    "class Conv2dBlock(nn.Module):\n",
    "    def __init__(self, input_dim, output_dim, kernel_size, stride, padding=0,\n",
    "                 conv_padding=0, dilation=1, weight_norm='none', norm='none',\n",
    "                 activation='relu', pad_type='zero', transpose=False):\n",
    "        \n",
    "        super(Conv2dBlock, self).__init__()\n",
    "        self.use_bias = True\n",
    "        # initialize padding\n",
    "        if pad_type == 'reflect':\n",
    "            self.pad = nn.ReflectionPad2d(padding)\n",
    "        elif pad_type == 'replicate':\n",
    "            self.pad = nn.ReplicationPad2d(padding)\n",
    "        elif pad_type == 'zero':\n",
    "            self.pad = nn.ZeroPad2d(padding)\n",
    "        elif pad_type == 'none':\n",
    "            self.pad = None\n",
    "        else:\n",
    "            assert 0, \"Unsupported padding type: {}\".format(pad_type)\n",
    "\n",
    "        # initialize normalization\n",
    "        norm_dim = output_dim\n",
    "        if norm == 'bn':\n",
    "            self.norm = nn.BatchNorm2d(norm_dim)\n",
    "        elif norm == 'in':\n",
    "            self.norm = nn.InstanceNorm2d(norm_dim)\n",
    "        elif norm == 'none':\n",
    "            self.norm = None\n",
    "        else:\n",
    "            assert 0, \"Unsupported normalization: {}\".format(norm)\n",
    "\n",
    "        if weight_norm == 'sn':\n",
    "            self.weight_norm = spectral_norm_fn\n",
    "        elif weight_norm == 'wn':\n",
    "            self.weight_norm = weight_norm_fn\n",
    "        elif weight_norm == 'none':\n",
    "            self.weight_norm = None\n",
    "        else:\n",
    "            assert 0, \"Unsupported normalization: {}\".format(weight_norm)\n",
    "\n",
    "        # initialize activation\n",
    "        if activation == 'relu':\n",
    "            self.activation = nn.ReLU(inplace=True)\n",
    "        elif activation == 'elu':\n",
    "            self.activation = nn.ELU(inplace=True)\n",
    "        elif activation == 'lrelu':\n",
    "            self.activation = nn.LeakyReLU(0.2, inplace=True)\n",
    "        elif activation == 'prelu':\n",
    "            self.activation = nn.PReLU()\n",
    "        elif activation == 'selu':\n",
    "            self.activation = nn.SELU(inplace=True)\n",
    "        elif activation == 'tanh':\n",
    "            self.activation = nn.Tanh()\n",
    "        elif activation == 'none':\n",
    "            self.activation = None\n",
    "        else:\n",
    "            assert 0, \"Unsupported activation: {}\".format(activation)\n",
    "\n",
    "        # initialize convolution\n",
    "        if transpose:\n",
    "            self.conv = nn.ConvTranspose2d(input_dim, output_dim,\n",
    "                                           kernel_size, stride,\n",
    "                                           padding=conv_padding,\n",
    "                                           output_padding=conv_padding,\n",
    "                                           dilation=dilation,\n",
    "                                           bias=self.use_bias)\n",
    "        else:\n",
    "            self.conv = nn.Conv2d(input_dim, output_dim, kernel_size, stride,\n",
    "                                  padding=conv_padding, dilation=dilation,\n",
    "                                  bias=self.use_bias)\n",
    "\n",
    "        if self.weight_norm:\n",
    "            self.conv = self.weight_norm(self.conv)\n",
    "\n",
    "    def forward(self, x):\n",
    "        if self.pad:\n",
    "            x = self.conv(self.pad(x))\n",
    "        else:\n",
    "            x = self.conv(x)\n",
    "        if self.norm:\n",
    "            x = self.norm(x)\n",
    "        if self.activation:\n",
    "            x = self.activation(x)\n",
    "        return x\n",
    "    \n",
    "    \n",
    "    \n",
    "def gen_conv(input_dim, output_dim, kernel_size=3, stride=1, padding=0, rate=1,\n",
    "             activation='elu', gated=False):\n",
    "    if gated:\n",
    "        conv2 = Conv2dBlockGated(input_dim, output_dim, kernel_size, stride,\n",
    "                           conv_padding=padding, dilation=rate,\n",
    "                           activation=activation)\n",
    "    else:\n",
    "        conv2 = Conv2dBlock(input_dim, output_dim, kernel_size, stride,\n",
    "                       conv_padding=padding, dilation=rate,\n",
    "                       activation=activation)\n",
    "        \n",
    "    return conv2\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "irish-relevance",
   "metadata": {},
   "outputs": [],
   "source": [
    "input_dim = config_gated['netG']['input_dim']\n",
    "gated = config_gated['netG']['gated']\n",
    "cnum = config_gated['netG']['ngf']\n",
    "\n",
    "conv1 = gen_conv(input_dim + 2, cnum, 5, 1, 2, gated=True).cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "inclusive-harbor",
   "metadata": {},
   "outputs": [],
   "source": [
    "_x1 = conv1(_in)\n",
    "print(_x1.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "robust-intervention",
   "metadata": {},
   "outputs": [],
   "source": [
    "test = torch.ones([4, 32, 256, 256])\n",
    "a, b = torch.chunk(test, 2, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "genuine-doctor",
   "metadata": {},
   "outputs": [],
   "source": [
    "c = a * b\n",
    "print(\"c: \", c.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "understanding-audit",
   "metadata": {},
   "outputs": [],
   "source": [
    "b.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "considered-integral",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_pytorch_latest_p36",
   "language": "python",
   "name": "conda_pytorch_latest_p36"
  },
  "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}