fastMRI_GAN.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "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.3"
    },
    "colab": {
      "name": "fastMRI_GAN.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/Nilanshrajput/FastMRI-Challenge/blob/master/fastMRI_GAN.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FOCbo8yYE4G1",
        "colab_type": "text"
      },
      "source": [
        "#### This notebook shows how to read the fastMRI dataset and apply some simple transformations to the data."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9_8uADLbE4G5",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "%matplotlib inline\n",
        "\n",
        "import h5py\n",
        "import numpy as np\n",
        "from matplotlib import pyplot as plt"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7wnaAGUZFJg-",
        "colab_type": "code",
        "outputId": "1b7696fd-ad47-4a8b-c04c-3351d90f5563",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 72
        }
      },
      "source": [
        "!curl -C - \"https://fastmri-dataset.s3.amazonaws.com/singlecoil_val.tar.gz?AWSAccessKeyId=AKIAJM2LEZ67Y2JL3KRA&Signature=b8SRruNRKvi0uJi4m7pWm9WmA7w%3D&Expires=1575704522\" --output singlecoil_val.tar.gz "
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "  % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current\n",
            "                                 Dload  Upload   Total   Spent    Left  Speed\n",
            "100 18.0G  100 18.0G    0     0  57.6M      0  0:05:20  0:05:20 --:--:-- 58.3M\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-VXO48AUFJla",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!tar xf singlecoil_val.tar.gz"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rVSZVFK1i1sP",
        "colab_type": "code",
        "outputId": "d9ef3e2e-edd2-4be8-f98c-d6ec9523cf52",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 72
        }
      },
      "source": [
        "!curl -C - \"https://fastmri-dataset.s3.amazonaws.com/singlecoil_train.tar.gz?AWSAccessKeyId=AKIAJM2LEZ67Y2JL3KRA&Signature=4pDyzc0pzFj0mMqXYgg3KOFpmiQ%3D&Expires=1575704522\" --output singlecoil_train.tar.gz"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "  % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current\n",
            "                                 Dload  Upload   Total   Spent    Left  Speed\n",
            "100 87.6G  100 87.6G    0     0  44.7M      0  0:33:24  0:33:24 --:--:-- 69.8M\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "m_q3FRE5i6D6",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!tar xf singlecoil_train.tar.gz"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KN4j03CAasiD",
        "colab_type": "code",
        "outputId": "6d612243-839c-4944-a63f-18c860fe8353",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 128
        }
      },
      "source": [
        "!pip install tensorboardX"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Requirement already satisfied: tensorboardX in /usr/local/lib/python3.6/dist-packages (1.8)\n",
            "Requirement already satisfied: protobuf>=3.2.0 in /usr/local/lib/python3.6/dist-packages (from tensorboardX) (3.7.1)\n",
            "Requirement already satisfied: six in /usr/local/lib/python3.6/dist-packages (from tensorboardX) (1.12.0)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.6/dist-packages (from tensorboardX) (1.16.5)\n",
            "Requirement already satisfied: setuptools in /usr/local/lib/python3.6/dist-packages (from protobuf>=3.2.0->tensorboardX) (41.2.0)\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gTsIZKlvOcQX",
        "colab_type": "code",
        "outputId": "ff546e00-a7e6-4982-846a-33c64b70ebc5",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 126
        }
      },
      "source": [
        "!git clone https://github.com/facebookresearch/fastMRI.git"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Cloning into 'fastMRI'...\n",
            "remote: Enumerating objects: 13, done.\u001b[K\n",
            "remote: Counting objects: 100% (13/13), done.\u001b[K\n",
            "remote: Compressing objects: 100% (11/11), done.\u001b[K\n",
            "remote: Total 93 (delta 3), reused 6 (delta 2), pack-reused 80\u001b[K\n",
            "Unpacking objects: 100% (93/93), done.\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "o7axuWk2atWr",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import numpy as np\n",
        "import torch\n",
        "import os\n",
        "from collections import OrderedDict\n",
        "from torch.autograd import Variable\n",
        "import torch.nn as nn\n",
        "# from torch.nn import init\n",
        "import functools\n",
        "import torchvision.models as models"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mfISHRZNh5l6",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from fastMRI.data import transforms\n",
        "from fastMRI.data.mri_data import SliceData\n",
        "from fastMRI.common.subsample import MaskFunc\n",
        "from torch.utils.data import DataLoader"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SX197k__PmLu",
        "colab_type": "code",
        "outputId": "22eeb6c7-bcaa-4553-9a03-3f02b41a517e",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 128
        }
      },
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Go to this URL in a browser: https://accounts.google.com/o/oauth2/auth?client_id=947318989803-6bn6qk8qdgf4n4g3pfee6491hc0brc4i.apps.googleusercontent.com&redirect_uri=urn%3Aietf%3Awg%3Aoauth%3A2.0%3Aoob&scope=email%20https%3A%2F%2Fwww.googleapis.com%2Fauth%2Fdocs.test%20https%3A%2F%2Fwww.googleapis.com%2Fauth%2Fdrive%20https%3A%2F%2Fwww.googleapis.com%2Fauth%2Fdrive.photos.readonly%20https%3A%2F%2Fwww.googleapis.com%2Fauth%2Fpeopleapi.readonly&response_type=code\n",
            "\n",
            "Enter your authorization code:\n",
            "··········\n",
            "Mounted at /content/drive\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NoiCv9KSQjDA",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "path_drive=pathlib.Path(\"/content/drive/My Drive/dl_projects/fast_mri/\")\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "x3GlCqKRQ4Z3",
        "colab_type": "code",
        "outputId": "87ce18ee-2e8c-479f-a517-416d6866dfaa",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 171
        }
      },
      "source": [
        "ls \"/content/drive/My Drive/dl_projects/fast_mri/\""
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "error",
          "ename": "AttributeError",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-49-2d9c75fbd01f>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mos\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mls\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpath_drive\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
            "\u001b[0;31mAttributeError\u001b[0m: module 'os' has no attribute 'ls'"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EuhhHGhJYfZu",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class ResnetGenerator(nn.Module):\n",
        "    def __init__(\n",
        "            self, input_nc, output_nc, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False,\n",
        "            n_blocks=6, gpu_ids=[], use_parallel=True, learn_residual=False, padding_type='reflect'):\n",
        "        assert (n_blocks >= 0)\n",
        "        super(ResnetGenerator, self).__init__()\n",
        "        self.input_nc = input_nc\n",
        "        self.output_nc = output_nc\n",
        "        self.ngf = ngf\n",
        "        self.gpu_ids = gpu_ids\n",
        "        self.use_parallel = use_parallel\n",
        "        self.learn_residual = learn_residual\n",
        "\n",
        "        if type(norm_layer) == functools.partial:\n",
        "            use_bias = norm_layer.func == nn.InstanceNorm2d\n",
        "        else:\n",
        "            use_bias = norm_layer == nn.InstanceNorm2d\n",
        "\n",
        "        model = [\n",
        "            nn.ReflectionPad2d(3),\n",
        "            nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0, bias=use_bias),\n",
        "            norm_layer(ngf),\n",
        "            nn.ReLU(True)\n",
        "        ]\n",
        "\n",
        "        n_downsampling = 2\n",
        "\n",
        "\n",
        "        model += [\n",
        "            nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1, bias=use_bias),\n",
        "            norm_layer(128),\n",
        "            nn.ReLU(True),\n",
        "\n",
        "            nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1, bias=use_bias),\n",
        "            norm_layer(256),\n",
        "            nn.ReLU(True)\n",
        "        ]\n",
        "\n",
        "        for i in range(n_blocks):\n",
        "           \n",
        "            model += [\n",
        "                ResnetBlock(256, padding_type=padding_type, norm_layer=norm_layer, use_dropout=use_dropout, use_bias=use_bias)\n",
        "            ]\n",
        "\n",
        "        model += [\n",
        "            nn.ConvTranspose2d(256, 128, kernel_size=3, stride=2, padding=1, output_padding=1, bias=use_bias),\n",
        "            norm_layer(128),\n",
        "            nn.ReLU(True),\n",
        "\n",
        "            nn.ConvTranspose2d(128, 64, kernel_size=3, stride=2, padding=1, output_padding=1, bias=use_bias),\n",
        "            norm_layer(64),\n",
        "            nn.ReLU(True),\n",
        "        ]\n",
        "\n",
        "        model += [\n",
        "            nn.ReflectionPad2d(3),\n",
        "            nn.Conv2d(64, output_nc, kernel_size=7, padding=0),\n",
        "            nn.Tanh()\n",
        "        ]\n",
        "\n",
        "        self.model = nn.Sequential(*model)\n",
        "\n",
        "    def forward(self, input):\n",
        "        if self.gpu_ids and isinstance(input.data, torch.cuda.FloatTensor) and self.use_parallel:\n",
        "            output = nn.parallel.data_parallel(self.model, input, self.gpu_ids)\n",
        "        else:\n",
        "            output = self.model(input)\n",
        "        if self.learn_residual:\n",
        "            # output = input + output\n",
        "            output = torch.clamp(input + output, min=-1, max=1)\n",
        "        return output\n",
        "\n",
        "\n",
        "# Define a resnet block\n",
        "class ResnetBlock(nn.Module):\n",
        "\n",
        "\tdef __init__(self, dim, padding_type, norm_layer, use_dropout, use_bias):\n",
        "\t\tsuper(ResnetBlock, self).__init__()\n",
        "\n",
        "\t\tpadAndConv = {\n",
        "\t\t\t'reflect': [\n",
        "                nn.ReflectionPad2d(1),\n",
        "                nn.Conv2d(dim, dim, kernel_size=3, bias=use_bias)],\n",
        "\t\t\t'replicate': [\n",
        "                nn.ReplicationPad2d(1),\n",
        "                nn.Conv2d(dim, dim, kernel_size=3, bias=use_bias)],\n",
        "\t\t\t'zero': [\n",
        "                nn.Conv2d(dim, dim, kernel_size=3, padding=1, bias=use_bias)]\n",
        "\t\t}\n",
        "\n",
        "\t\ttry:\n",
        "\t\t\tblocks = padAndConv[padding_type] + [\n",
        "\t\t\t\tnorm_layer(dim),\n",
        "\t\t\t\tnn.ReLU(True)\n",
        "            ] + [\n",
        "\t\t\t\tnn.Dropout(0.5)\n",
        "\t\t\t] if use_dropout else [] + padAndConv[padding_type] + [\n",
        "\t\t\t\tnorm_layer(dim)\n",
        "\t\t\t]\n",
        "\t\texcept:\n",
        "\t\t\traise NotImplementedError('padding [%s] is not implemented' % padding_type)\n",
        "\n",
        "\t\tself.conv_block = nn.Sequential(*blocks)\n",
        "\n",
        "\n",
        "\tdef forward(self, x):\n",
        "\t\tout = x + self.conv_block(x)\n",
        "\t\treturn out\n",
        "\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ClT1KDJ6YtR8",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Defines the PatchGAN discriminator with the specified arguments.\n",
        "class NLayerDiscriminator(nn.Module):\n",
        "    def __init__(self, input_nc, ndf=64, n_layers=10, norm_layer=nn.BatchNorm2d, use_sigmoid=False, gpu_ids=[],\n",
        "                 use_parallel=True):\n",
        "        super(NLayerDiscriminator, self).__init__()\n",
        "        self.gpu_ids = gpu_ids\n",
        "        self.use_parallel = use_parallel\n",
        "\n",
        "        if type(norm_layer) == functools.partial:\n",
        "            use_bias = norm_layer.func == nn.InstanceNorm2d\n",
        "        else:\n",
        "            use_bias = norm_layer == nn.InstanceNorm2d\n",
        "\n",
        "        kw = 4\n",
        "        padw = int(np.ceil((kw - 1) / 2))\n",
        "        sequence = [\n",
        "            nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw),\n",
        "            nn.LeakyReLU(0.2, True)\n",
        "        ]\n",
        "\n",
        "        nf_mult = 1\n",
        "        nf_mult_prev = 1\n",
        "        for n in range(1, n_layers):\n",
        "            nf_mult_prev = nf_mult\n",
        "            nf_mult = min(2 ** n, 8)\n",
        "            sequence += [\n",
        "                nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult,\n",
        "                          kernel_size=kw, stride=2, padding=padw, bias=use_bias),\n",
        "                norm_layer(ndf * nf_mult),\n",
        "                nn.LeakyReLU(0.2, True)\n",
        "            ]\n",
        "\n",
        "        nf_mult_prev = nf_mult\n",
        "        nf_mult = min(2 ** n_layers, 8)\n",
        "        sequence += [\n",
        "            nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=1, padding=padw, bias=use_bias),\n",
        "            norm_layer(ndf * nf_mult),\n",
        "            nn.LeakyReLU(0.2, True)\n",
        "        ]\n",
        "\n",
        "        sequence += [nn.Conv2d(ndf * nf_mult, 1, kernel_size=kw, stride=1, padding=padw)]\n",
        "\n",
        "        if use_sigmoid:\n",
        "            sequence += [nn.Sigmoid()]\n",
        "\n",
        "        self.model = nn.Sequential(*sequence)\n",
        "\n",
        "    def forward(self, input):\n",
        "        if len(self.gpu_ids) and isinstance(input.data, torch.cuda.FloatTensor) and self.use_parallel:\n",
        "            return nn.parallel.data_parallel(self.model, input, self.gpu_ids)\n",
        "        else:\n",
        "            return self.model(input)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Y6g-ON6lBpkU",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class ImagePool():\n",
        "    def __init__(self, pool_size):\n",
        "        self.pool_size = pool_size\n",
        "        if self.pool_size > 0:\n",
        "            self.num_imgs = 0\n",
        "            self.images = []\n",
        "\n",
        "    def query(self, images):\n",
        "        if self.pool_size == 0:\n",
        "            return images\n",
        "        return_images = []\n",
        "        for image in images.data:\n",
        "            image = torch.unsqueeze(image, 0)\n",
        "            if self.num_imgs < self.pool_size:\n",
        "                self.num_imgs = self.num_imgs + 1\n",
        "                self.images.append(image)\n",
        "                return_images.append(image)\n",
        "            else:\n",
        "                p = random.uniform(0, 1)\n",
        "                if p > 0.5:\n",
        "                    random_id = random.randint(0, self.pool_size-1)\n",
        "                    tmp = self.images[random_id].clone()\n",
        "                    self.images[random_id] = image\n",
        "                    return_images.append(tmp)\n",
        "                else:\n",
        "                    return_images.append(image)\n",
        "        return_images = Variable(torch.cat(return_images, 0))\n",
        "        return return_images\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2xCHBqZtYvs8",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class ContentLoss:\n",
        "\tdef __init__(self, loss):\n",
        "\t\tself.criterion = loss\n",
        "\t\t\t\n",
        "\tdef get_loss(self, fakeIm, realIm):\n",
        "\t\treturn self.criterion(fakeIm, realIm)\n",
        "\n",
        "class PerceptualLoss():\n",
        "\t\n",
        "  def contentFunc(self):\n",
        "    conv_3_3_layer = 14\n",
        "    cnn = models.vgg19(pretrained=True).features\n",
        "    cnn = cnn.cuda()\n",
        "    model = nn.Sequential()\n",
        "    model = model.cuda()\n",
        "    for i,layer in enumerate(list(cnn)):\n",
        "      model.add_module(str(i),layer)\n",
        "      if i == conv_3_3_layer:\n",
        "        break\n",
        "    return model\n",
        "\n",
        "  def __init__(self, loss):\n",
        "    self.criterion = loss\n",
        "    self.contentFunc = self.contentFunc()\n",
        "\n",
        "  def get_loss(self, fakeIm, realIm):\n",
        "    new_fakeIm = fakeIm[:,:, :, :] * torch.ones(3)[None,:, None, None].to(torch.device('cuda') )\n",
        "\n",
        "    new_realIm = realIm[:,:, :, :] * torch.ones(3)[None,:, None, None].to(torch.device('cuda') )\n",
        "\n",
        "\n",
        "    f_fake = self.contentFunc.forward(new_fakeIm)\n",
        "\n",
        "\n",
        "    f_real = self.contentFunc.forward(new_realIm)\n",
        "    f_real_no_grad = f_real\n",
        "    loss = self.criterion(f_fake, f_real_no_grad)\n",
        "    return loss\n",
        "\n",
        "class GANLoss(nn.Module):\n",
        "\tdef __init__(\n",
        "\t\t\tself, use_l1=True, target_real_label=1.0,\n",
        "\t\t\ttarget_fake_label=0.0, tensor=torch.FloatTensor):\n",
        "\t\tsuper(GANLoss, self).__init__()\n",
        "\t\tself.real_label = target_real_label\n",
        "\t\tself.fake_label = target_fake_label\n",
        "\t\tself.real_label_var = None\n",
        "\t\tself.fake_label_var = None\n",
        "\t\tself.Tensor = tensor\n",
        "\t\tif use_l1:\n",
        "\t\t\tself.loss = nn.L1Loss()\n",
        "\t\telse:\n",
        "\t\t\tself.loss = nn.BCELoss()\n",
        "\n",
        "\tdef get_target_tensor(self, input, target_is_real):\n",
        "\t\ttarget_tensor = None\n",
        "\t\tif target_is_real:\n",
        "\t\t\tcreate_label = ((self.real_label_var is None) or\n",
        "\t\t\t\t\t\t\t(self.real_label_var.numel() != input.numel()))\n",
        "\t\t\tif create_label:\n",
        "\t\t\t\treal_tensor = self.Tensor(input.size()).fill_(self.real_label)\n",
        "\t\t\t\tself.real_label_var = Variable(real_tensor, requires_grad=False)\n",
        "\t\t\ttarget_tensor = self.real_label_var\n",
        "\t\telse:\n",
        "\t\t\tcreate_label = ((self.fake_label_var is None) or\n",
        "\t\t\t\t\t\t\t(self.fake_label_var.numel() != input.numel()))\n",
        "\t\t\tif create_label:\n",
        "\t\t\t\tfake_tensor = self.Tensor(input.size()).fill_(self.fake_label)\n",
        "\t\t\t\tself.fake_label_var = Variable(fake_tensor, requires_grad=False)\n",
        "\t\t\ttarget_tensor = self.fake_label_var\n",
        "\t\treturn target_tensor\n",
        "\n",
        "\tdef __call__(self, input, target_is_real):\n",
        "\t\ttarget_tensor = self.get_target_tensor(input, target_is_real).to(torch.device('cuda'))\n",
        "\t\treturn self.loss(input, target_tensor)\n",
        "    \n",
        "\n",
        "class DiscLoss:\n",
        "\tdef name(self):\n",
        "\t\treturn 'DiscLoss'\n",
        "\n",
        "\tdef __init__(self):\n",
        "\t\tself.criterionGAN = GANLoss(use_l1=False)\n",
        "\t\tself.fake_AB_pool = ImagePool(pool_size=50)\n",
        "\t\t\n",
        "\tdef get_g_loss(self,net, realA, fakeB):\n",
        "\t\t# First, G(A) should fake the discriminator\n",
        "\t\tpred_fake = net.forward(fakeB)\n",
        "\t\treturn self.criterionGAN(pred_fake, 1)\n",
        "\t\t\n",
        "\tdef get_loss(self, net, realA, fakeB, realB):\n",
        "\t\t# Fake\n",
        "\t\t# stop backprop to the generator by detaching fake_B\n",
        "\t\t# Generated Image Disc Output should be close to zero\n",
        "\t\tself.pred_fake = net.forward(fakeB)\n",
        "\t\tself.loss_D_fake = self.criterionGAN(self.pred_fake, 0)\n",
        "\n",
        "\t\t# Real\n",
        "\t\tself.pred_real = net.forward(realB)\n",
        "\t\tself.loss_D_real = self.criterionGAN(self.pred_real, 1)\n",
        "\n",
        "\t\t# Combined loss\n",
        "\t\tself.loss_D = (self.loss_D_fake + self.loss_D_real) * 0.5\n",
        "\t\treturn self.loss_D\n",
        "\t\t\n",
        "class DiscLossLS(DiscLoss):\n",
        "\tdef name(self):\n",
        "\t\treturn 'DiscLossLS'\n",
        "\n",
        "\tdef __init__(self):\n",
        "\t\tsuper(DiscLoss, self).__init__()\n",
        "\t\t# DiscLoss.initialize(self, opt, tensor)\n",
        "\t\tself.criterionGAN = GANLoss(use_l1=True, tensor=tensor)\n",
        "\t\t\n",
        "\tdef get_g_loss(self,net, realA, fakeB):\n",
        "\t\treturn DiscLoss.get_g_loss(self,net, realA, fakeB)\n",
        "\t\t\n",
        "\tdef get_loss(self, net, realA, fakeB, realB):\n",
        "\t\treturn DiscLoss.get_loss(self, net, realA, fakeB, realB)\n",
        "\n",
        "        "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ikGD2rfnY0Vm",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "\n",
        "\n",
        "class DiscLossWGANGP(DiscLossLS):\n",
        "    def name(self):\n",
        "      return 'DiscLossWGAN-GP'\n",
        "\n",
        "    def __init__():\n",
        "      super(DiscLossWGANGP, self).__init__()\n",
        "      # DiscLossLS.initialize(self, opt, tensor)\n",
        "      self.LAMBDA = 10\n",
        "        \n",
        "\n",
        "\n",
        "\n",
        "    def get_g_loss(net, realA, fakeB):\n",
        "        # First, G(A) should fake the discriminator\n",
        "        D_fake = net.forward(fakeB)\n",
        "        return -D_fake.mean()\n",
        "\n",
        "    def calc_gradient_penalty(self, netD, real_data, fake_data):\n",
        "        alpha = torch.rand(1, 1)\n",
        "        alpha = alpha.expand(real_data.size())\n",
        "        alpha = alpha.cuda()\n",
        "\n",
        "        interpolates = alpha * real_data + ((1 - alpha) * fake_data)\n",
        "\n",
        "        interpolates = interpolates.cuda()\n",
        "        interpolates = Variable(interpolates, requires_grad=True)\n",
        "\n",
        "        disc_interpolates = netD.forward(interpolates)\n",
        "\n",
        "        gradients = autograd.grad(\n",
        "            outputs=disc_interpolates, inputs=interpolates, grad_outputs=torch.ones(disc_interpolates.size()).cuda(),\n",
        "            create_graph=True, retain_graph=True, only_inputs=True\n",
        "        )[0]\n",
        "\n",
        "        gradient_penalty = ((gradients.norm(2, dim=1) - 1) ** 2).mean() * self.LAMBDA\n",
        "        return gradient_penalty\n",
        "\n",
        "    def get_loss( net, realA, fakeB, realB):\n",
        "        D_fake = net.forward(fakeB)\n",
        "        D_fake = D_fake.mean()\n",
        "\n",
        "        # Real\n",
        "        D_real = net.forward(realB)\n",
        "        D_real = D_real.mean()\n",
        "        # Combined loss\n",
        "        loss_D = D_fake - D_real\n",
        "        gradient_penalty = calc_gradient_penalty(net, realB.data, fakeB.data)\n",
        "        return loss_D + gradient_penalty\n",
        "\n",
        "\t\t"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0adDk39PcOpD",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def weights_init(m):\n",
        "    classname = m.__class__.__name__\n",
        "    if classname.find('Conv') != -1:\n",
        "        m.weight.data.normal_(0.0, 0.02)\n",
        "        if hasattr(m.bias, 'data'):\n",
        "            m.bias.data.fill_(0)\n",
        "    elif classname.find('BatchNorm2d') != -1:\n",
        "        m.weight.data.normal_(1.0, 0.02)\n",
        "        m.bias.data.fill_(0)\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CzBbEkU0bN3c",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "channel_rate = 64\n",
        "# Note the image_shape must be multiple of patch_shape\n",
        "image_shape = (320, 320, 1)\n",
        "patch_shape = (channel_rate, channel_rate, 3)\n",
        "\n",
        "\n",
        "ngf = 64\n",
        "ndf = 64\n",
        "input_nc = 1\n",
        "output_nc = 1\n",
        "\n",
        "use_dropout=True\n",
        "gpu_ids=[0]\n",
        "use_parallel=False\n",
        "learn_residual=False\n",
        "use_sigmoid= True\n",
        "\n",
        "norm_layer = functools.partial(nn.InstanceNorm2d, affine=False, track_running_stats=True)\n",
        "\n",
        "netG = ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=12,\n",
        "                               gpu_ids=gpu_ids, use_parallel=use_parallel, learn_residual=learn_residual)\n",
        "\n",
        "netG.cuda(gpu_ids[0])\n",
        "netG.apply(weights_init)\n",
        "\n",
        "n_layers_D=12\n",
        "\n",
        "netD = NLayerDiscriminator(input_nc, ndf, n_layers_D, norm_layer=norm_layer, use_sigmoid=use_sigmoid,\n",
        "                                   gpu_ids=gpu_ids, use_parallel=use_parallel)\n",
        "\n",
        "netD.cuda(gpu_ids[0])\n",
        "netD.apply(weights_init)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7uqyjy-Rha7k",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class DataTransform:\n",
        "    def __init__(self, mask_func, resolution, which_challenge, use_seed=True):\n",
        "        self.resolution = resolution\n",
        "        self.mask_func = mask_func\n",
        "    def __call__(self, kspace, target, attrs, fname, slice):\n",
        "        kspace = transforms.to_tensor(kspace)\n",
        "        \n",
        "        seed = tuple(map(ord, fname))\n",
        "        masked_kspace, _ = transforms.apply_mask(kspace, self.mask_func, seed)\n",
        "\n",
        "      \n",
        "        image = transforms.ifft2(masked_kspace)\n",
        "        image = transforms.complex_center_crop(image, (self.resolution, self.resolution))\n",
        "        image = transforms.complex_abs(image)\n",
        "        image, mean, std = transforms.normalize_instance(image)\n",
        "        image = image.clamp(-6, 6)\n",
        "        target = transforms.to_tensor(target)\n",
        "        # Normalize target\n",
        "        target = transforms.normalize(target, mean, std, eps=1e-11)\n",
        "        target = target.clamp(-6, 6)\n",
        "        return image, target, mean, std\n",
        "        #return image, mean, std, fname, slice"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LKqpJDoDhfb7",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import pathlib\n",
        "\n",
        "\n",
        "\n",
        "root = pathlib.Path('/content/singlecoil_train')\n",
        "root_val=pathlib.Path('/content/singlecoil_val')\n",
        "batch_size=1\n",
        "num_workers=8\n",
        "\n",
        "def create_dataset(root, resolution=320, center_fractions=[0.08], accelerations=[4]):\n",
        "    \n",
        "    mask_func = MaskFunc(center_fractions, accelerations)\n",
        "    data = SliceData(\n",
        "        root = root,\n",
        "        transform = DataTransform(mask_func, resolution,'singlecoil'),\n",
        "        sample_rate = 1.0,\n",
        "        challenge = 'singlecoil'\n",
        "    )\n",
        "    return data\n",
        "\n",
        "data_train=create_dataset(root)\n",
        "data_val=create_dataset(root_val)\n",
        "    \n",
        "def create_dataloader(data, batch_size, num_workers):\n",
        "    dataloader = DataLoader(\n",
        "        dataset = data,\n",
        "        batch_size = batch_size,\n",
        "        num_workers = num_workers,\n",
        "    )\n",
        "    return dataloader\n",
        "\n",
        "\n",
        "data_load_train=create_dataloader(data_train,batch_size, num_workers)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yD7uwldnP_QV",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "beta1=0.5\n",
        "lr=0.0001\n",
        "device = torch.device('cuda') "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kslyzuTdY2Mb",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "\n",
        "optimizer_G = torch.optim.Adam( netG.parameters(), lr=lr, betas=(beta1, 0.999) )\n",
        "optimizer_D = torch.optim.Adam( netD.parameters(), lr=lr, betas=(beta1, 0.999) )\n",
        "\n",
        "def forward(netG,input_A,input_B):\n",
        "\t\treal_A = Variable(input_A)\n",
        "\t\tfake_B = netG.forward(real_A)\n",
        "\t\treal_B = Variable(input_B)\n",
        "\n",
        "\n",
        "def backward_D(netD, real_A, fake_B, real_B, eval=False):\n",
        "  discLoss=DiscLoss()\n",
        "  loss_D = discLoss.get_loss(netD, real_A, fake_B, real_B)\n",
        "\n",
        "  if not eval:\n",
        "    #retaingraph=true save the gradients of prev batch(its in use as we are\n",
        "    #calling backward() multiple times on same batch multiple times look optimize_parameters func )\n",
        "    loss_D.backward(retain_graph=True)\n",
        "  \n",
        "  return loss_D.item()\n",
        "  \n",
        "def backward_G(netD, real_A, fake_B, real_B,eval=False):\n",
        "    discLoss=DiscLoss()\n",
        "    per_loss=PerceptualLoss(nn.MSELoss())\n",
        "    contentLoss=  ContentLoss(nn.L1Loss())\n",
        "    #print(fake_B.shape)\n",
        "    #print(real_A.shape)\n",
        "    loss_G_GAN = discLoss.get_g_loss(netD, real_A, fake_B)\n",
        "    # Second, G(A) = B\n",
        "    loss_G_Content = contentLoss.get_loss(fake_B, real_B) *100 \n",
        "    loss_G_percept= per_loss.get_loss(fake_B, real_B)\n",
        "\n",
        "    loss_G = loss_G_GAN + loss_G_Content +loss_G_percept*50\n",
        "    \n",
        "    if not eval:\n",
        "      loss_G.backward()\n",
        "    \n",
        "    return loss_G.item()\n",
        "\n",
        "    \n",
        "def optimize_parameters(netD,netG, real_A, fake_B, real_B,optimizer_G,optimizer_D):\n",
        "    forward(netG,real_A,real_B)\n",
        "    criticUpdates=5\n",
        "    lossd=[]\n",
        "    for iter_d in range(criticUpdates):\n",
        "        optimizer_D.zero_grad()\n",
        "        lossd.append(backward_D(netD, real_A, fake_B, real_B))\n",
        "        optimizer_D.step()\n",
        "    loss_D=torch.mean(lossd)\n",
        "    optimizer_G.zero_grad()\n",
        "    loss_G= backward_G(netD, real_A, fake_B, real_B)\n",
        "    optimizer_G.step()\n",
        "    \n",
        "    return loss_G,loss_D"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "aaysHyPwQIrd",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "os.mkdir(path_drive/'weights')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WjxmpoUFSrPP",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "\n",
        "weights_dir=path_drive/'weights'"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Zc2Icg4edHO8",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "#stores the weights in  weights folder\n",
        "def save_network(network, network_label, epoch_label, gpu_ids):\n",
        "  save_filename = '%s_net_%s.pth' % (epoch_label, network_label)\n",
        "  save_path = os.path.join(weights_dir, save_filename)\n",
        "  torch.save(network.cpu().state_dict(), save_path)\n",
        "  if len(gpu_ids) and torch.cuda.is_available():\n",
        "      network.cuda(device=gpu_ids[0])"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HESHa6v1eQCq",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# run this before model initialization\n",
        "def load_network(network, network_label, epoch_label):\n",
        "  save_filename = '%s_net_%s.pth' % (epoch_label, network_label)\n",
        "  save_path = os.path.join(weights_dir, save_filename)\n",
        "  network.load_state_dict(torch.load(save_path))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dP0jfz9fh6St",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def generate_images(netG, test_input, tar):\n",
        "  # the training=True is intentional here since\n",
        "  # we want the batch statistics while running the netG\n",
        "  # on the test dataset. If we use training=False, we will get\n",
        "  # the accumulated statistics learned from the training dataset\n",
        "  # (which we don't want)\n",
        "  netG.eval()\n",
        "  prediction = netG[0].forward(test_input)\n",
        "  plt.figure(figsize=(15,15))\n",
        "\n",
        "  display_list = [test_input[0], tar[0], prediction[0]]\n",
        "  title = ['Input Image', 'Ground Truth', 'Predicted Image']\n",
        "\n",
        "  for i in range(3):\n",
        "    plt.subplot(1, 3, i+1)\n",
        "    plt.title(title[i])\n",
        "    # getting the pixel values between [0, 1] to plot it.\n",
        "    plt.imshow(display_list[i] * 0.5 + 0.5)\n",
        "    plt.axis('off')\n",
        "  plt.show()\n",
        "  netG.train()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XsgBerpKMJEq",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "report_interval=100"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OiLtlhrCNTIe",
        "colab_type": "code",
        "outputId": "0d441dfa-7ec0-42d4-984c-508df70cf690",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 135
        }
      },
      "source": [
        "def evaluate( epoch, netG, netD, data_loader,opt_g, opt_d, writer):\n",
        "    netG.eval()\n",
        "    losses = []\n",
        "    start = time.perf_counter()\n",
        "    with torch.no_grad():\n",
        "        for iter, data in enumerate(data_loader):\n",
        "            input, target, mean, std, norm = data\n",
        "            input = input.unsqueeze(1).to(device)\n",
        "            target = target.to(device)\n",
        "            output = netG(input).squeeze(1)\n",
        "\n",
        "            mean = mean.unsqueeze(1).unsqueeze(2).to(device)\n",
        "            std = std.unsqueeze(1).unsqueeze(2).to(device)\n",
        "#            target = target * std + mean\n",
        "#            output = output * std + mean\n",
        "\n",
        "#            norm = norm.unsqueeze(1).unsqueeze(2).to(args.device)\n",
        "            loss = bakward_G(netD, input, output, target, eval= True)\n",
        "            losses.append(loss.item())\n",
        "        writer.add_scalar('Dev_Loss', np.mean(losses), epoch)\n",
        "    return np.mean(losses), time.perf_counter() - start"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "error",
          "ename": "SyntaxError",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;36m  File \u001b[0;32m\"<ipython-input-57-67b022099788>\"\u001b[0;36m, line \u001b[0;32m1\u001b[0m\n\u001b[0;31m    def evaluate( epoch, netG, data_loader,,opt_g, opt_d, writer):\u001b[0m\n\u001b[0m                                           ^\u001b[0m\n\u001b[0;31mSyntaxError\u001b[0m\u001b[0;31m:\u001b[0m invalid syntax\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6kRO7BnKZfST",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "device = torch.device('cuda') \n",
        "epoch_count=5\n",
        "batchSize=1\n",
        "\n",
        "\n",
        "import time\n",
        "def train(data_loader, netG,netD, writer):\n",
        "  #dataset = data_loader.load_data()\n",
        "  dataset_size = len(data_loader)\n",
        "  print('#training images = %d' % dataset_size)\n",
        "  total_steps = 0\n",
        "  \n",
        "  train_time=0\n",
        "  for epoch in range(epoch_count):\n",
        "    epoch_start_time = time.time()\n",
        "    epoch_iter = 0\n",
        "    start_epoch = start_iter = time.perf_counter()\n",
        "    global_step = epoch * len(data_loader)\n",
        "    for i, data in enumerate(data_loader):\n",
        "      iter_start_time = time.time()\n",
        "      total_steps += batchSize\n",
        "      epoch_iter += batchSize\n",
        "      input, target, mean, std = data\n",
        "      #print(input.unsqueeze(1).shape)\n",
        "      #input-input.expand(batchSize,320,320,1)\n",
        "      #target=target.expand(batchSize,320,320,1)\n",
        "      input = input.unsqueeze(1).to(device)\n",
        "      #print(input.shape)\n",
        "      target = target.unsqueeze(1).to(device)\n",
        "      output = netG(input)#.squeeze(1)\n",
        "      #print(target.shape)\n",
        "      loss_g,loss_d=optimize_parameters(netD,netG,input,output,target,optimizer_G,optimizer_D)\n",
        "      \n",
        "      \n",
        "      TrainLoss=loss_g+loss_d\n",
        "      writer.add_scalar('TrainLoss',TrainLoss,global_step+i)\n",
        "      avg_loss = 0.99 * avg_loss + 0.01 * TrainLoss if i > 0 else TrainLoss.item()\n",
        "      writer.add_scalar('TrainLoss', loss.item(), global_step + iter)\n",
        "\n",
        "      if iter % report_interval == 0:\n",
        "          logging.info(\n",
        "              f'Epoch = [{epoch:3d}/{epoch_count:3d}] '\n",
        "              f'Iter = [{i:4d}/{len(data_loader):4d}] '\n",
        "              f'Loss = {TrainLoss:.4g} Avg Loss = {avg_loss:.4g} '\n",
        "              f'Time = {time.perf_counter() - start_iter:.4f}s',\n",
        "          )\n",
        "        \n",
        "      start_iter = time.perf_counter()\n",
        "      train_time= time.perf_counter() - start_epoch\n",
        "      \n",
        "      \n",
        "      \n",
        "      if total_steps % save_latest_freq == 0:\n",
        "\t\t\t\tprint('saving the latest model (epoch %d, total_steps %d)' % (epoch, total_steps))\n",
        "\t\t\t\tsave_network(netG, 'G', 'latest', gpu_ids) #gpu_ids daal dena\n",
        "\t\t    save_network(netD, 'D', 'latest', gpu_ids)\n",
        "        print('Visualising images :- ')\n",
        "        generate_images(netG, input, target)\n",
        "      \n",
        "    if epoch % save_epoch_freq == 0:\n",
        "      print('saving the model at the end of epoch %d, iters %d' % (epoch, total_steps))\n",
        "      save_network(netG, 'G', 'latest', gpu_ids)\n",
        "\t\t  save_network(netD, 'D', 'latest', gpu_ids)\n",
        "      save_network(netG, 'G', epoch, gpu_ids)\n",
        "\t\t  save_network(netD, 'D', epoch, gpu_ids)\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "A17GZtmNnCXd",
        "colab_type": "code",
        "outputId": "c346d5cd-a874-4f14-8e02-16cfcbc6ff6e",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        }
      },
      "source": [
        "train(data_load_train,netG,netD)"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "#training images = 34742\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n",
            "torch.Size([1, 1, 320, 320])\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "error",
          "ename": "KeyboardInterrupt",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-26-60c26a945fc1>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mtrain\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdata_load_train\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mnetG\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mnetD\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
            "\u001b[0;32m<ipython-input-24-3902a7333bfa>\u001b[0m in \u001b[0;36mtrain\u001b[0;34m(data_loader, netG, netD)\u001b[0m\n\u001b[1;32m     26\u001b[0m       \u001b[0moutput\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mnetG\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;31m#.squeeze(1)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     27\u001b[0m       \u001b[0;31m#print(target.shape)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 28\u001b[0;31m       \u001b[0moptimize_parameters\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnetD\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mnetG\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0moutput\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mtarget\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0moptimizer_G\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0moptimizer_D\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
            "\u001b[0;32m<ipython-input-23-5edbe9003f6d>\u001b[0m in \u001b[0;36moptimize_parameters\u001b[0;34m(netD, netG, real_A, fake_B, real_B, optimizer_G, optimizer_D)\u001b[0m\n\u001b[1;32m     38\u001b[0m     \u001b[0;32mfor\u001b[0m \u001b[0miter_d\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcriticUpdates\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     39\u001b[0m         \u001b[0moptimizer_D\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mzero_grad\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 40\u001b[0;31m         \u001b[0mbackward_D\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnetD\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mreal_A\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfake_B\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mreal_B\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     41\u001b[0m         \u001b[0moptimizer_D\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstep\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     42\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m<ipython-input-23-5edbe9003f6d>\u001b[0m in \u001b[0;36mbackward_D\u001b[0;34m(netD, real_A, fake_B, real_B)\u001b[0m\n\u001b[1;32m     12\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mbackward_D\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnetD\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mreal_A\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfake_B\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mreal_B\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     13\u001b[0m   \u001b[0mdiscLoss\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mDiscLoss\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 14\u001b[0;31m   \u001b[0mloss_D\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mdiscLoss\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mget_loss\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnetD\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mreal_A\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfake_B\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mreal_B\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     15\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     16\u001b[0m   \u001b[0mloss_D\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mretain_graph\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m<ipython-input-6-7b4afd2307df>\u001b[0m in \u001b[0;36mget_loss\u001b[0;34m(self, net, realA, fakeB, realB)\u001b[0m\n\u001b[1;32m     87\u001b[0m                 \u001b[0;31m# Generated Image Disc Output should be close to zero\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     88\u001b[0m                 \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpred_fake\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mnet\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfakeB\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 89\u001b[0;31m                 \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mloss_D_fake\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcriterionGAN\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpred_fake\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     90\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     91\u001b[0m                 \u001b[0;31m# Real\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m<ipython-input-6-7b4afd2307df>\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, input, target_is_real)\u001b[0m\n\u001b[1;32m     65\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     66\u001b[0m         \u001b[0;32mdef\u001b[0m \u001b[0m__call__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtarget_is_real\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 67\u001b[0;31m                 \u001b[0mtarget_tensor\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mget_target_tensor\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtarget_is_real\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'cuda'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     68\u001b[0m                 \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mloss\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtarget_tensor\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     69\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vrSVK1qeGM9B",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def generate_images(netG, test_input, tar):\n",
        "  # the training=True is intentional here since\n",
        "  # we want the batch statistics while running the netG\n",
        "  # on the test dataset. If we use training=False, we will get\n",
        "  # the accumulated statistics learned from the training dataset\n",
        "  # (which we don't want)\n",
        "  prediction = netG(test_input, training=True)\n",
        "  plt.figure(figsize=(15,15))\n",
        "\n",
        "  display_list = [test_input[0], tar[0], prediction[0]]\n",
        "  title = ['Input Image', 'Ground Truth', 'Predicted Image']\n",
        "\n",
        "  for i in range(3):\n",
        "    plt.subplot(1, 3, i+1)\n",
        "    plt.title(title[i])\n",
        "    # getting the pixel values between [0, 1] to plot it.\n",
        "    plt.imshow(display_list[i] * 0.5 + 0.5)\n",
        "    plt.axis('off')\n",
        "  plt.show()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ry7Fly2XzWq2",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SJgpKF7S3AiY",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}