{"cells":[{"metadata":{},"cell_type":"markdown","source":"\n"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport os\nimport matplotlib.pyplot as plt\n\nimport torch\nfrom torch import nn, optim\nimport torch.nn.functional as F\nfrom torchvision import datasets, transforms\nfrom torchvision.utils import save_image\nfrom torch.utils.data import Dataset, DataLoader\nfrom PIL import Image\n\nfrom tqdm import tqdm_notebook as tqdm","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Generator and Discriminator"},{"metadata":{"trusted":true},"cell_type":"code","source":"class Generator(nn.Module):\n    def __init__(self, nz=128, channels=3):\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=3):\n        super(Discriminator, self).__init__()\n        \n        self.channels = channels\n\n        def convlayer(n_input, n_output, k_size=3, 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":{},"cell_type":"markdown","source":"## Parameters of GAN"},{"metadata":{"trusted":true},"cell_type":"code","source":"batch_size = 32\nlr = 0.003\nbeta1 = 0.5\nepochs = 300\n\nreal_label = 0.5\nfake_label = 0\nnz = 128\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Data loader and Augmentations"},{"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        self.img_names = os.listdir(img_dir)\n        self.transform1 = transform1\n        self.transform2 = transform2\n        \n        self.imgs = []\n        for img_name in self.img_names:\n            img = Image.open(os.path.join(img_dir, img_name))\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":{"_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","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/all-dogs/all-dogs/',\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":{},"cell_type":"markdown","source":"### Let's see examples of traing data"},{"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":{},"cell_type":"markdown","source":"## Initialize models and optimizers"},{"metadata":{"trusted":true},"cell_type":"code","source":"netG = Generator(nz).to(device)\nnetD = Discriminator().to(device)\n\ncriterion = nn.BCELoss()\n\noptimizerD = optim.Adam(netD.parameters(), lr=lr, betas=(beta1, 0.999))\noptimizerG = optim.Adam(netG.parameters(), lr=lr, betas=(beta1, 0.999))\n\nfixed_noise = torch.randn(25, nz, 1, 1, device=device)","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":{},"cell_type":"markdown","source":"## Training loop"},{"metadata":{"trusted":true},"cell_type":"code","source":"for 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        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()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Let's visualise generated results "},{"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)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Make predictions and submit"},{"metadata":{"trusted":true},"cell_type":"code","source":"if not os.path.exists('../output_images'):\n    os.mkdir('../output_images')\nim_batch_size = 50\nn_images=10000\nfor i_batch in range(0, n_images, im_batch_size):\n    gen_z = torch.randn(im_batch_size, nz, 1, 1, device=device)\n    gen_images = (netG(gen_z) + 1)/2\n    images = gen_images.to(\"cpu\").clone().detach()\n    images = images.numpy().transpose(0, 2, 3, 1)\n    for i_image in range(gen_images.size(0)):\n        save_image(gen_images[i_image, :, :, :], os.path.join('../output_images', f'image_{i_batch+i_image:05d}.png'))\n\n\nimport shutil\nshutil.make_archive('images', 'zip', '../output_images')","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.6.4","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"}},"nbformat":4,"nbformat_minor":1}