{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Challenge: Generative Adversarial Network\n### Generative Dog Images"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import os\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# torch\nimport torch\nfrom torch import nn, optim\nimport torch.nn.functional as F\nimport torchvision.utils as utils\nfrom torch.utils.data import Dataset, DataLoader\nfrom torchvision import datasets, transforms\nfrom torchvision.utils import save_image\nfrom torch.autograd import Variable\nfrom torchvision.models.vgg import vgg16\nfrom torchvision.transforms import Compose, RandomCrop, ToTensor, ToPILImage, CenterCrop, Resize\n\nimport argparse\nimport time\nfrom math import log10\nfrom PIL import Image\nfrom tqdm import tqdm_notebook as tqdm\nimport xml.etree.ElementTree as ET # for parsing XML\nimport math\nfrom math import exp\n\nimport matplotlib\nimport matplotlib.pyplot as plt # plotting\n\nimport shutil # used for making zip archives of output\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class Generator(nn.Module):\n    def __init__(self, nz, channels):\n        super(Generator, self).__init__()\n        \n        self.nz = nz\n        self.channels = channels\n        \n        def convlayer(n_input, n_output, k_size=4, stride=2, padding=0):\n            block = [\n                nn.ConvTranspose2d(n_input, n_output, kernel_size=k_size, stride=stride, padding=padding, bias=False),\n                nn.BatchNorm2d(n_output),\n                nn.ReLU(inplace=True),\n            ]\n            return block\n\n        self.model = nn.Sequential(\n            *convlayer(self.nz, 1024, 4, 1, 0), # Fully connected layer via convolution.\n            *convlayer(1024, 512, 4, 2, 1),\n            *convlayer(512, 256, 4, 2, 1),\n            *convlayer(256, 128, 4, 2, 1),\n            *convlayer(128, 64, 4, 2, 1),\n            nn.ConvTranspose2d(64, self.channels, 3, 1, 1),\n            nn.Tanh()\n        )\n\n\n    def forward(self, z):\n        z = z.view(-1, self.nz, 1, 1)\n        img = self.model(z)\n        return img\n\nclass Discriminator(nn.Module):\n    def __init__(self, channels):\n        super(Discriminator, self).__init__()\n        \n        self.channels = channels\n\n        def convlayer(n_input, n_output, k_size=4, stride=2, padding=0, bn=False):\n            block = [nn.Conv2d(n_input, n_output, kernel_size=k_size, stride=stride, padding=padding, bias=False)]\n            if bn:\n                block.append(nn.BatchNorm2d(n_output))\n            block.append(nn.LeakyReLU(0.2, inplace=True))\n            return block\n\n        self.model = nn.Sequential(\n            *convlayer(self.channels, 32, 4, 2, 1),\n            *convlayer(32, 64, 4, 2, 1),\n            *convlayer(64, 128, 4, 2, 1, bn=True),\n            *convlayer(128, 256, 4, 2, 1, bn=True),\n            nn.Conv2d(256, 1, 4, 1, 0, bias=False),  # FC with Conv.\n        )\n\n    def forward(self, imgs):\n        out = self.model(imgs)\n        return out.view(-1, 1)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"batch_size = 32\nlr = 0.0003\nbeta1 = 0.5\n\nnz = 128\nfixed_noise = torch.randn(25, nz, 1, 1, device=device)\n\nreal_label = 0.9\n\nepochs = 200\n\ncriterion = nn.BCELoss()\n\nnetG = Generator(nz, 3).to(device)\nnetD = Discriminator(3).to(device)\n\noptimizerD = optim.Adam(netD.parameters(), lr=lr, betas=(beta1, 0.999))\noptimizerG = optim.Adam(netG.parameters(), lr=lr, betas=(beta1, 0.999))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def bounding_box(image):\n    image = '../input/stanford-dogs/annotation/Annotation/' + os.path.splitext(image)[0]\n    tree = ET.parse(image)\n    root = tree.getroot()\n    objects = root.findall('object')\n    for o in objects:\n        bndbox = o.find('bndbox') # reading bound box\n        xmin = int(bndbox.find('xmin').text)\n        ymin = int(bndbox.find('ymin').text)\n        xmax = int(bndbox.find('xmax').text)\n        ymax = int(bndbox.find('ymax').text)\n        \n    return (xmin,ymin,xmax,ymax)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class DogDataset(Dataset):\n    def __init__(self, img_dir, transform1=None, transform2=None):\n    \n        self.img_dir = img_dir\n        img_directory_names = os.listdir(img_dir)\n        self.img_names = []\n        for img_directory_name in img_directory_names:\n            list_img_directory_files = os.listdir(img_dir + \"/\" + img_directory_name)\n            for list_img_directory_file in list_img_directory_files:    \n                self.img_names.append(img_directory_name + \"/\" + list_img_directory_file)\n                        \n        self.transform1 = transform1\n        self.transform2 = transform2\n        \n        self.imgs = []\n        for img_name in self.img_names:\n            #bbox=bounding_box(img_name)\n            img = Image.open(os.path.join(img_dir, img_name)).convert('RGB')\n            #img = img.crop(bbox)\n            \n            if self.transform1 is not None:\n                img = self.transform1(img)\n                \n            self.imgs.append(img)\n\n    def __getitem__(self, index):\n        img = self.imgs[index]\n        \n        if self.transform2 is not None:\n            img = self.transform2(img)\n        \n        return img\n\n    def __len__(self):\n        return len(self.imgs)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"transform1 = transforms.Compose([transforms.Resize(64),\n                                 transforms.CenterCrop(64)])\n\n# Data augmentation and converting to tensors\nrandom_transforms = [transforms.RandomRotation(degrees=5)]\ntransform2 = transforms.Compose([transforms.RandomHorizontalFlip(p=0.5),\n                                 transforms.RandomApply(random_transforms, p=0.3), \n                                 transforms.ToTensor(),\n                                 transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])\n                                 \ntrain_dataset = DogDataset(img_dir='../input/stanford-dogs/images/Images',\n                           transform1=transform1,\n                           transform2=transform2)\n\ntrain_loader = DataLoader(dataset=train_dataset,\n                          batch_size=batch_size,\n                          shuffle=True,\n                          num_workers=4)\n\n                                           \nimgs = next(iter(train_loader))\nimgs = imgs.numpy().transpose(0, 2, 3, 1)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"fig = plt.figure(figsize=(25, 16))\nfor ii, img in enumerate(imgs):\n    ax = fig.add_subplot(4, 8, ii + 1, xticks=[], yticks=[])\n    \n    plt.imshow((img+1)/2)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def show_generated_img():\n    noise = torch.randn(1, nz, 1, 1, device=device)\n    gen_image = netG(noise).to(\"cpu\").clone().detach().squeeze(0)\n    gen_image = gen_image.numpy().transpose(1, 2, 0)\n    plt.imshow((gen_image+1)/2)\n    plt.show()","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"G_losses = []\nD_losses = []\n\nfor epoch in range(epochs):\n    \n    for ii, real_images in tqdm(enumerate(train_loader), total=len(train_loader)):\n        ############################\n        # (1) Update D network\n        ###########################\n        netD.zero_grad()\n        real_images = real_images.to(device)\n        batch_size = real_images.size(0)\n        labels = torch.full((batch_size, 1), real_label, device=device)\n        outputR = netD(real_images)\n        noise = torch.randn(batch_size, nz, 1, 1, device=device)\n        fake = netG(noise)\n        outputF = netD(fake.detach())\n        errD = (torch.mean((outputR - torch.mean(outputF) - labels) ** 2) + \n                torch.mean((outputF - torch.mean(outputR) + labels) ** 2))/2\n        errD.backward(retain_graph=True)\n        optimizerD.step()\n        ############################\n        # (2) Update G network\n        ###########################\n        netG.zero_grad()\n        outputF = netD(fake)   \n        errG = (torch.mean((outputR - torch.mean(outputF) + labels) ** 2) +\n                torch.mean((outputF - torch.mean(outputR) - labels) ** 2))/2\n        errG.backward()\n        optimizerG.step()\n        \n        # Save Losses for plotting later\n        G_losses.append(errG.item())\n        D_losses.append(errD.item())\n        \n        if (ii+1) % (len(train_loader)//2) == 0:\n            print('[%d/%d][%d/%d] Loss_D: %.4f Loss_G: %.4f'\n                  % (epoch + 1, epochs, ii+1, len(train_loader),\n                     errD.item(), errG.item()))\n\n    show_generated_img()\n        \ntorch.save(netG.state_dict(), 'generator.pth')\ntorch.save(netD.state_dict(), 'discriminator.pth')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plt.figure(figsize=(10,5))\nplt.title(\"Generator and Discriminator Loss During Training\")\nplt.plot(G_losses,label=\"Generator\")\nplt.plot(D_losses,label=\"Discriminator\")\nplt.xlabel(\"Iterations\")\nplt.ylabel(\"Loss\")\nplt.legend()\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"gen_z = torch.randn(32, nz, 1, 1, device=device)\ngen_images = (netG(gen_z).to(\"cpu\").clone().detach() + 1)/2\ngen_images = gen_images.numpy().transpose(0, 2, 3, 1)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"fig = plt.figure(figsize=(25, 16))\nfor ii, img in enumerate(gen_images):\n    ax = fig.add_subplot(4, 8, ii + 1, xticks=[], yticks=[])\n    plt.imshow(img)\n    matplotlib.image.imsave(str(ii) + '_gen_img.png', img)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class SRGenerator(nn.Module):\n    def __init__(self, scale_factor):\n        upsample_block_num = int(math.log(scale_factor, 2))\n\n        super(SRGenerator, self).__init__()\n        self.block1 = nn.Sequential(\n            nn.Conv2d(3, 64, kernel_size=9, padding=4),\n            nn.PReLU()\n        )\n        self.block2 = ResidualBlock(64)\n        self.block3 = ResidualBlock(64)\n        self.block4 = ResidualBlock(64)\n        self.block5 = ResidualBlock(64)\n        self.block6 = ResidualBlock(64)\n        self.block7 = nn.Sequential(\n            nn.Conv2d(64, 64, kernel_size=3, padding=1),\n            nn.BatchNorm2d(64)\n        )\n        block8 = [UpsampleBLock(64, 2) for _ in range(upsample_block_num)]\n        block8.append(nn.Conv2d(64, 3, kernel_size=9, padding=4))\n        self.block8 = nn.Sequential(*block8)\n\n    def forward(self, x):\n        block1 = self.block1(x)\n        block2 = self.block2(block1)\n        block3 = self.block3(block2)\n        block4 = self.block4(block3)\n        block5 = self.block5(block4)\n        block6 = self.block6(block5)\n        block7 = self.block7(block6)\n        block8 = self.block8(block1 + block7)\n\n        return (torch.tanh(block8) + 1) / 2\n\n\nclass SRDiscriminator(nn.Module):\n    def __init__(self):\n        super(SRDiscriminator, self).__init__()\n        self.net = nn.Sequential(\n            nn.Conv2d(3, 64, kernel_size=3, padding=1),\n            nn.LeakyReLU(0.2),\n\n            nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=1),\n            nn.BatchNorm2d(64),\n            nn.LeakyReLU(0.2),\n\n            nn.Conv2d(64, 128, kernel_size=3, padding=1),\n            nn.BatchNorm2d(128),\n            nn.LeakyReLU(0.2),\n\n            nn.Conv2d(128, 128, kernel_size=3, stride=2, padding=1),\n            nn.BatchNorm2d(128),\n            nn.LeakyReLU(0.2),\n\n            nn.Conv2d(128, 256, kernel_size=3, padding=1),\n            nn.BatchNorm2d(256),\n            nn.LeakyReLU(0.2),\n\n            nn.Conv2d(256, 256, kernel_size=3, stride=2, padding=1),\n            nn.BatchNorm2d(256),\n            nn.LeakyReLU(0.2),\n\n            nn.Conv2d(256, 512, kernel_size=3, padding=1),\n            nn.BatchNorm2d(512),\n            nn.LeakyReLU(0.2),\n\n            nn.Conv2d(512, 512, kernel_size=3, stride=2, padding=1),\n            nn.BatchNorm2d(512),\n            nn.LeakyReLU(0.2),\n\n            nn.AdaptiveAvgPool2d(1),\n            nn.Conv2d(512, 1024, kernel_size=1),\n            nn.LeakyReLU(0.2),\n            nn.Conv2d(1024, 1, kernel_size=1)\n        )\n\n    def forward(self, x):\n        batch_size = x.size(0)\n        return torch.sigmoid(self.net(x).view(batch_size))\n\n\nclass ResidualBlock(nn.Module):\n    def __init__(self, channels):\n        super(ResidualBlock, self).__init__()\n        self.conv1 = nn.Conv2d(channels, channels, kernel_size=3, padding=1)\n        self.bn1 = nn.BatchNorm2d(channels)\n        self.prelu = nn.PReLU()\n        self.conv2 = nn.Conv2d(channels, channels, kernel_size=3, padding=1)\n        self.bn2 = nn.BatchNorm2d(channels)\n\n    def forward(self, x):\n        residual = self.conv1(x)\n        residual = self.bn1(residual)\n        residual = self.prelu(residual)\n        residual = self.conv2(residual)\n        residual = self.bn2(residual)\n\n        return x + residual\n\n\nclass UpsampleBLock(nn.Module):\n    def __init__(self, in_channels, up_scale):\n        super(UpsampleBLock, self).__init__()\n        self.conv = nn.Conv2d(in_channels, in_channels * up_scale ** 2, kernel_size=3, padding=1)\n        self.pixel_shuffle = nn.PixelShuffle(up_scale)\n        self.prelu = nn.PReLU()\n\n    def forward(self, x):\n        x = self.conv(x)\n        x = self.pixel_shuffle(x)\n        x = self.prelu(x)\n        return x","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class GeneratorLoss(nn.Module):\n    def __init__(self):\n        super(GeneratorLoss, self).__init__()\n        vgg = vgg16(pretrained=True)\n        loss_network = nn.Sequential(*list(vgg.features)[:31]).eval()\n        for param in loss_network.parameters():\n            param.requires_grad = False\n        self.loss_network = loss_network\n        self.mse_loss = nn.MSELoss()\n        self.tv_loss = TVLoss()\n\n    def forward(self, out_labels, out_images, target_images):\n        # Adversarial Loss\n        adversarial_loss = torch.mean(1 - out_labels)\n        # Perception Loss\n        perception_loss = self.mse_loss(self.loss_network(out_images), self.loss_network(target_images))\n        # Image Loss\n        image_loss = self.mse_loss(out_images, target_images)\n        # TV Loss\n        tv_loss = self.tv_loss(out_images)\n        return image_loss + 0.001 * adversarial_loss + 0.006 * perception_loss + 2e-8 * tv_loss\n\n\nclass TVLoss(nn.Module):\n    def __init__(self, tv_loss_weight=1):\n        super(TVLoss, self).__init__()\n        self.tv_loss_weight = tv_loss_weight\n\n    def forward(self, x):\n        batch_size = x.size()[0]\n        h_x = x.size()[2]\n        w_x = x.size()[3]\n        count_h = self.tensor_size(x[:, :, 1:, :])\n        count_w = self.tensor_size(x[:, :, :, 1:])\n        h_tv = torch.pow((x[:, :, 1:, :] - x[:, :, :h_x - 1, :]), 2).sum()\n        w_tv = torch.pow((x[:, :, :, 1:] - x[:, :, :, :w_x - 1]), 2).sum()\n        return self.tv_loss_weight * 2 * (h_tv / count_h + w_tv / count_w) / batch_size\n\n    @staticmethod\n    def tensor_size(t):\n        return t.size()[1] * t.size()[2] * t.size()[3]\n\n\nif __name__ == \"__main__\":\n    g_loss = GeneratorLoss()\n    print(g_loss)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def is_image_file(filename):\n    return any(filename.endswith(extension) for extension in ['.png', '.jpg', '.jpeg', '.PNG', '.JPG', '.JPEG'])\n\n\ndef calculate_valid_crop_size(crop_size, upscale_factor):\n    return crop_size - (crop_size % upscale_factor)\n\n\ndef train_hr_transform(crop_size):\n    return Compose([\n        RandomCrop(crop_size),\n        ToTensor(),\n    ])\n\n\ndef train_lr_transform(crop_size, upscale_factor):\n    return Compose([\n        ToPILImage(),\n        Resize(crop_size // upscale_factor, interpolation=Image.BICUBIC),\n        ToTensor()\n    ])\n\n\ndef display_transform():\n    return Compose([\n        ToPILImage(),\n        Resize(400),\n        CenterCrop(400),\n        ToTensor()\n    ])\n\n\nclass TrainDatasetFromFolder(Dataset):\n    def __init__(self, dataset_dir, crop_size, upscale_factor):\n        super(TrainDatasetFromFolder, self).__init__()\n        self.image_filenames = [os.path.join(dataset_dir, x) for x in os.listdir(dataset_dir) if is_image_file(x)]\n        crop_size = calculate_valid_crop_size(crop_size, upscale_factor)\n        self.hr_transform = train_hr_transform(crop_size)\n        self.lr_transform = train_lr_transform(crop_size, upscale_factor)\n\n    def __getitem__(self, index):\n        hr_image = self.hr_transform(Image.open(self.image_filenames[index]))\n        lr_image = self.lr_transform(hr_image)\n        return lr_image, hr_image\n\n    def __len__(self):\n        return len(self.image_filenames)\n\n\nclass ValDatasetFromFolder(Dataset):\n    def __init__(self, dataset_dir, upscale_factor):\n        super(ValDatasetFromFolder, self).__init__()\n        self.upscale_factor = upscale_factor\n        self.image_filenames = [os.path.join(dataset_dir, x) for x in os.listdir(dataset_dir) if is_image_file(x)]\n\n    def __getitem__(self, index):\n        hr_image = Image.open(self.image_filenames[index])\n        w, h = hr_image.size\n        crop_size = calculate_valid_crop_size(min(w, h), self.upscale_factor)\n        lr_scale = Resize(crop_size // self.upscale_factor, interpolation=Image.BICUBIC)\n        hr_scale = Resize(crop_size, interpolation=Image.BICUBIC)\n        hr_image = CenterCrop(crop_size)(hr_image)\n        lr_image = lr_scale(hr_image)\n        hr_restore_img = hr_scale(lr_image)\n        return ToTensor()(lr_image), ToTensor()(hr_restore_img), ToTensor()(hr_image)\n\n    def __len__(self):\n        return len(self.image_filenames)\n\n\nclass TestDatasetFromFolder(Dataset):\n    def __init__(self, dataset_dir, upscale_factor):\n        super(TestDatasetFromFolder, self).__init__()\n        self.lr_path = dataset_dir + '/SRF_' + str(upscale_factor) + '/data/'\n        self.hr_path = dataset_dir + '/SRF_' + str(upscale_factor) + '/target/'\n        self.upscale_factor = upscale_factor\n        self.lr_filenames = [os.path.join(self.lr_path, x) for x in os.listdir(self.lr_path) if is_image_file(x)]\n        self.hr_filenames = [os.path.join(self.hr_path, x) for x in os.listdir(self.hr_path) if is_image_file(x)]\n\n    def __getitem__(self, index):\n        image_name = self.lr_filenames[index].split('/')[-1]\n        lr_image = Image.open(self.lr_filenames[index])\n        w, h = lr_image.size\n        hr_image = Image.open(self.hr_filenames[index])\n        hr_scale = Resize((self.upscale_factor * h, self.upscale_factor * w), interpolation=Image.BICUBIC)\n        hr_restore_img = hr_scale(lr_image)\n        return image_name, ToTensor()(lr_image), ToTensor()(hr_restore_img), ToTensor()(hr_image)\n\n    def __len__(self):\n        return len(self.lr_filenames)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def ssim(img1, img2, window_size = 11, size_average = True):\n    (_, channel, _, _) = img1.size()\n    window = create_window(window_size, channel)\n    \n    if img1.is_cuda:\n        window = window.cuda(img1.get_device())\n    window = window.type_as(img1)\n    \n    return _ssim(img1, img2, window, window_size, channel, size_average)\n\ndef gaussian(window_size, sigma):\n    gauss = torch.Tensor([exp(-(x - window_size//2)**2/float(2*sigma**2)) for x in range(window_size)])\n    return gauss/gauss.sum()\n\ndef create_window(window_size, channel):\n    _1D_window = gaussian(window_size, 1.5).unsqueeze(1)\n    _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0)\n    window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous())\n    return window\n\ndef _ssim(img1, img2, window, window_size, channel, size_average = True):\n    mu1 = F.conv2d(img1, window, padding = window_size//2, groups = channel)\n    mu2 = F.conv2d(img2, window, padding = window_size//2, groups = channel)\n\n    mu1_sq = mu1.pow(2)\n    mu2_sq = mu2.pow(2)\n    mu1_mu2 = mu1*mu2\n\n    sigma1_sq = F.conv2d(img1*img1, window, padding = window_size//2, groups = channel) - mu1_sq\n    sigma2_sq = F.conv2d(img2*img2, window, padding = window_size//2, groups = channel) - mu2_sq\n    sigma12 = F.conv2d(img1*img2, window, padding = window_size//2, groups = channel) - mu1_mu2\n\n    C1 = 0.01**2\n    C2 = 0.03**2\n\n    ssim_map = ((2*mu1_mu2 + C1)*(2*sigma12 + C2))/((mu1_sq + mu2_sq + C1)*(sigma1_sq + sigma2_sq + C2))\n\n    if size_average:\n        return ssim_map.mean()\n    else:\n        return ssim_map.mean(1).mean(1).mean(1)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"if __name__ == '__main__':    \n    CROP_SIZE = 88\n    UPSCALE_FACTOR = 4\n    NUM_EPOCHS = 100\n    \n    train_set = TrainDatasetFromFolder('../input/div2k-train/DIV2K_train_HR', crop_size=CROP_SIZE, upscale_factor=UPSCALE_FACTOR)\n    val_set = ValDatasetFromFolder('../input/div2k-valid/DIV2K_valid_HR', upscale_factor=UPSCALE_FACTOR)\n    train_loader = DataLoader(dataset=train_set, num_workers=4, batch_size=64, shuffle=True)\n    val_loader = DataLoader(dataset=val_set, num_workers=4, batch_size=1, shuffle=False)\n    \n    netG = SRGenerator(UPSCALE_FACTOR)\n    print('# generator parameters:', sum(param.numel() for param in netG.parameters()))\n    netD = SRDiscriminator()\n    print('# discriminator parameters:', sum(param.numel() for param in netD.parameters()))\n    \n    generator_criterion = GeneratorLoss()\n    \n    if torch.cuda.is_available():\n        netG.cuda()\n        netD.cuda()\n        generator_criterion.cuda()\n    \n    optimizerG = optim.Adam(netG.parameters())\n    optimizerD = optim.Adam(netD.parameters())\n    \n    results = {'d_loss': [], 'g_loss': [], 'd_score': [], 'g_score': [], 'psnr': [], 'ssim': []}\n    \n    for epoch in range(1, NUM_EPOCHS + 1):\n        train_bar = tqdm(train_loader)\n        running_results = {'batch_sizes': 0, 'd_loss': 0, 'g_loss': 0, 'd_score': 0, 'g_score': 0}\n    \n        netG.train()\n        netD.train()\n        for data, target in train_bar:\n            g_update_first = True\n            batch_size = data.size(0)\n            running_results['batch_sizes'] += batch_size\n    \n            ############################\n            # (1) Update D network: maximize D(x)-1-D(G(z))\n            ###########################\n            real_img = Variable(target)\n            if torch.cuda.is_available():\n                real_img = real_img.cuda()\n            z = Variable(data)\n            if torch.cuda.is_available():\n                z = z.cuda()\n            fake_img = netG(z)\n    \n            netD.zero_grad()\n            real_out = netD(real_img).mean()\n            fake_out = netD(fake_img).mean()\n            d_loss = 1 - real_out + fake_out\n            d_loss.backward(retain_graph=True)\n            optimizerD.step()\n    \n            ############################\n            # (2) Update G network: minimize 1-D(G(z)) + Perception Loss + Image Loss + TV Loss\n            ###########################\n            netG.zero_grad()\n            g_loss = generator_criterion(fake_out, fake_img, real_img)\n            g_loss.backward()\n            \n            fake_img = netG(z)\n            fake_out = netD(fake_img).mean()\n            \n            \n            optimizerG.step()\n\n            # loss for current batch before optimization \n            running_results['g_loss'] += g_loss.item() * batch_size\n            running_results['d_loss'] += d_loss.item() * batch_size\n            running_results['d_score'] += real_out.item() * batch_size\n            running_results['g_score'] += fake_out.item() * batch_size\n    \n            train_bar.set_description(desc='[%d/%d] Loss_D: %.4f Loss_G: %.4f D(x): %.4f D(G(z)): %.4f' % (\n                epoch, NUM_EPOCHS, running_results['d_loss'] / running_results['batch_sizes'],\n                running_results['g_loss'] / running_results['batch_sizes'],\n                running_results['d_score'] / running_results['batch_sizes'],\n                running_results['g_score'] / running_results['batch_sizes']))\n    \n        netG.eval()\n        out_path = 'training_results/SRF_' + str(UPSCALE_FACTOR) + '/'\n        if not os.path.exists(out_path):\n            os.makedirs(out_path)\n        \n        with torch.no_grad():\n            val_bar = tqdm(val_loader)\n            valing_results = {'mse': 0, 'ssims': 0, 'psnr': 0, 'ssim': 0, 'batch_sizes': 0}\n            val_images = []\n            for val_lr, val_hr_restore, val_hr in val_bar:\n                batch_size = val_lr.size(0)\n                valing_results['batch_sizes'] += batch_size\n                lr = val_lr\n                hr = val_hr\n                if torch.cuda.is_available():\n                    lr = lr.cuda()\n                    hr = hr.cuda()\n                sr = netG(lr)\n        \n                batch_mse = ((sr - hr) ** 2).data.mean()\n                valing_results['mse'] += batch_mse * batch_size\n                batch_ssim = ssim(sr, hr).item()\n                valing_results['ssims'] += batch_ssim * batch_size\n                valing_results['psnr'] = 10 * log10(1 / (valing_results['mse'] / valing_results['batch_sizes']))\n                valing_results['ssim'] = valing_results['ssims'] / valing_results['batch_sizes']\n                val_bar.set_description(\n                    desc='[converting LR images to SR images] PSNR: %.4f dB SSIM: %.4f' % (\n                        valing_results['psnr'], valing_results['ssim']))\n        \n                val_images.extend(\n                    [display_transform()(val_hr_restore.squeeze(0)), display_transform()(hr.data.cpu().squeeze(0)),\n                     display_transform()(sr.data.cpu().squeeze(0))])\n            val_images = torch.stack(val_images)\n            val_images = torch.chunk(val_images, val_images.size(0) // 15)\n            val_save_bar = tqdm(val_images, desc='[saving training results]')\n            index = 1\n            for image in val_save_bar:\n                image = utils.make_grid(image, nrow=3, padding=5)\n                #utils.save_image(image, out_path + 'epoch_%d_index_%d.png' % (epoch, index), padding=5)\n                index += 1\n    \n        # save model parameters\n        torch.save(netG.state_dict(), 'netG_epoch_%d_%d.pth' % (UPSCALE_FACTOR, epoch))\n        torch.save(netD.state_dict(), 'netD_epoch_%d_%d.pth' % (UPSCALE_FACTOR, epoch))\n        # save loss\\scores\\psnr\\ssim\n        results['d_loss'].append(running_results['d_loss'] / running_results['batch_sizes'])\n        results['g_loss'].append(running_results['g_loss'] / running_results['batch_sizes'])\n        results['d_score'].append(running_results['d_score'] / running_results['batch_sizes'])\n        results['g_score'].append(running_results['g_score'] / running_results['batch_sizes'])\n        results['psnr'].append(valing_results['psnr'])\n        results['ssim'].append(valing_results['ssim'])\n    \n        #if epoch % 10 == 0 and epoch != 0:\n        #    out_path = 'statistics/'\n        #    data_frame = pd.DataFrame(\n        #        data={'Loss_D': results['d_loss'], 'Loss_G': results['g_loss'], 'Score_D': results['d_score'],\n        #              'Score_G': results['g_score'], 'PSNR': results['psnr'], 'SSIM': results['ssim']},\n        #        index=range(1, epoch + 1))\n        #    data_frame.to_csv(out_path + 'srf_' + str(UPSCALE_FACTOR) + '_train_results.csv', index_label='Epoch')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"for gen_img_nr in range(0, 31):\n    UPSCALE_FACTOR = 4\n    TEST_MODE = True\n    IMAGE_NAME = '../working/' + str(gen_img_nr) + '_gen_img.png'\n    MODEL_NAME = 'netG_epoch_4_100.pth'\n\n    model = SRGenerator(UPSCALE_FACTOR).eval()\n    if TEST_MODE:\n        model.cuda()\n        model.load_state_dict(torch.load(MODEL_NAME))\n    else:\n        model.load_state_dict(torch.load(MODEL_NAME, map_location=lambda storage, loc: storage))\n\n    image = Image.open(IMAGE_NAME).convert('RGB')\n    image = Variable(ToTensor()(image), volatile=True).unsqueeze(0)\n    if TEST_MODE:\n        image = image.cuda()\n\n    start = time.clock()\n    out = model(image)\n    elapsed = (time.clock() - start)\n    print('cost' + str(elapsed) + 's')\n    out_img = ToPILImage()(out[0].data.cpu())\n    plt.imshow(out_img)\n    out_img.save('out_srf_' + str(gen_img_nr) + '_' + \"test.png\")","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":1}