{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load in \n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\nimport torch\nfrom torch import nn, functional\nimport torchvision\nfrom skimage import io, transform, exposure, util\nfrom random import shuffle, seed, choice, random\nfrom collections.abc import Iterable\nfrom tqdm.auto import tqdm\nimport itertools as it\n\n# Plotting\nfrom plotly.offline import iplot, init_notebook_mode\nfrom plotly import graph_objs as go\ninit_notebook_mode(connected=True)\nfrom matplotlib import pyplot as plt\n\n# Input data files are available in the \"../input/\" directory.\n# For example, running this (by clicking run or pressing Shift+Enter) will list the files in the input directory\n\nimport os\nimport sys\n\nos.makedirs('./output', exist_ok=True)\n        \nprint(os.listdir(\"../input/\"))\n\n# Any results you write to the current directory are saved as output.","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class DogLoader(torch.utils.data.Dataset):\n    def __init__(self, path, preload_steps = 100, random_batch_size = 32, shuffle_images = True, augment=True, preload_all=False, shuffle_seed = None):\n        self.preload_all = preload_all\n        self.augment = augment\n        self.preload_steps = preload_steps\n        self.random_batch_size = random_batch_size\n        self.path = path\n        self.files = sorted(os.listdir(path))\n        self.file_indices = list(range(len(self.files)))\n        if shuffle_images:\n            if shuffle_seed is not None:\n                seed(shuffle_seed)\n            shuffle(self.files)\n        \n        if preload_all:\n            self.preload_dict = {}\n            for file in tqdm(self.files):\n                self.preload_dict[self.path+file] = io.imread(self.path+file)\n        \n        self.counter = 0\n        self.preload = self[0:self.preload_steps]\n        self.preload_counter = 0\n        \n    def __len__(self):\n        return len(self.files)\n        \n    def __iter__(self):\n        self.counter = 0\n        self.preload = self[0:self.preload_steps]\n        self.preload_counter = 0\n        return self\n    \n    def __next__(self):\n        if self.counter < len(self):\n            if self.preload_steps <= self.preload_counter:\n                self.preload = self[self.counter:min(self.counter+self.preload_steps, len(self))]\n                self.preload_counter = 0\n            self.preload_counter += 1\n            self.counter += 1\n            return self.preload[self.preload_counter-1]\n        else:\n            raise StopIteration\n                \n        raise StopIteration\n    \n    def __getitem__(self, idx):\n        if type(idx) is int:\n            return self[idx:idx+1]\n        if isinstance(idx, Iterable):\n            used_files = [self.files[i] for i in idx]\n        else:\n            used_files = self.files[idx]\n        images = []\n        for file in used_files:\n            if self.preload_all:\n                image = self.preload_dict[self.path+file]\n            else:\n                image = io.imread(self.path+file)\n            w,h,c = image.shape\n            if self.augment:\n                if random() > 0.5:\n                    image = image[:,::-1,:]\n                image = transform.rotate(image, -5+10*random(), resize=True)\n                image = image[w//2-min(w,h)//2:w//2+min(w,h)//2,h//2-min(w,h)//2:h//2+min(w,h)//2,:]\n                image = np.concatenate([exposure.equalize_hist(image[:,:,[i]]) for i in range(3)], axis=2)\n            image = transform.resize(image, (64,64))\n            images.append(image)\n        return torch.from_numpy(np.stack(images)).permute(0,3,1,2).float()\n    \n    def get_random_batch(self):\n        idx = [choice(self.file_indices) for i in range(self.random_batch_size)]\n        return self[idx]\n\ndef test_DogLoader(length = 30):\n    dogs1 = DogLoader(\"../input/all-dogs/all-dogs/\", preload_steps=3, augment=False, shuffle_images=True, shuffle_seed=1)\n    dogs2 = DogLoader(\"../input/all-dogs/all-dogs/\", preload_steps=5, augment=False, shuffle_images=True, shuffle_seed=1)\n    i = 0\n    for image1, image2 in zip(dogs1, dogs2):\n        print(f\"{dogs1.counter}:{dogs1.preload_counter} {dogs2.counter}:{dogs2.preload_counter}\")\n        assert (image1==image2).all()\n        i += 1\n        if i >= length:\n            break\n\ntest_DogLoader()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"data_loader = torch.utils.data.DataLoader(\n    dataset=DogLoader(path=\"../input/all-dogs/all-dogs/\"),\n    batch_size=32,\n    num_workers=4,\n    drop_last=True,\n    collate_fn=torch.cat\n)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class Discriminator(nn.Module):\n    def __init__(self):\n        super(Discriminator, self).__init__()\n        self.conv0 = nn.Conv2d(in_channels = 3, out_channels = 128, kernel_size = (4,4), stride = 1, padding = 0, bias = True)\n        self.bn0 = nn.BatchNorm2d(128)\n        self.act0 = nn.ReLU()\n        self.conv1 = nn.Conv2d(in_channels = 128, out_channels = 64, kernel_size = (4,4), stride = 2, padding = 0, bias = True)\n        self.bn1 = nn.BatchNorm2d(64)\n        self.act1 = nn.ReLU()\n        self.conv2 = nn.Conv2d(in_channels = 64, out_channels = 32, kernel_size = (4,4), stride = 2, padding = 0, bias = True)\n        self.bn2 = nn.BatchNorm2d(32)\n        self.act2 = nn.ReLU()\n        self.conv3 = nn.Conv2d(in_channels = 32, out_channels = 16, kernel_size = (4,4), stride = 2, padding = 0, bias = True)\n        self.bn3 = nn.BatchNorm2d(16)\n        self.act3 = nn.ReLU()\n        self.dense4 = nn.Linear(in_features=400, out_features=1)\n        self.act4 = nn.Sigmoid()\n        \n    def forward(self, x: torch.Tensor):\n        x = self.conv0(x)\n        x = self.bn0(x)\n        x = self.act0(x)\n        #print(x.shape)\n        x = self.conv1(x)\n        x = self.bn1(x)\n        x = self.act1(x)\n        #print(x.shape)\n        x = self.conv2(x)\n        x = self.bn2(x)\n        x = self.act2(x)\n        #print(x.shape)\n        x = self.conv3(x)\n        x = self.bn3(x)\n        x = self.act3(x)\n        #print(x.shape)\n        x = self.dense4(x.view(-1, 16*5*5))\n        x = self.act4(x)\n        #print(x.shape)\n        return x\n    \n    \nclass Generator(nn.Module):\n    def __init__(self):\n        super(Generator, self).__init__()\n        self.dense0 = nn.Linear(in_features = 6, out_features = 16*8*8)\n        self.act0 = nn.ReLU()\n        self.deconv1 = nn.ConvTranspose2d(in_channels=16, out_channels= 32, kernel_size=(4,4), stride = 2, padding = 0, bias = True)\n        self.bn1 = nn.BatchNorm2d(32)\n        self.act1 = nn.ReLU()\n        self.deconv2 = nn.ConvTranspose2d(in_channels=32, out_channels= 64, kernel_size=(4,4), stride = 2, padding = 0, bias = True)\n        self.bn2 = nn.BatchNorm2d(64)\n        self.act2 = nn.ReLU()\n        self.deconv3 = nn.ConvTranspose2d(in_channels=64, out_channels= 128, kernel_size=(4,4), stride = 2, padding = 0, bias = True)\n        self.bn3= nn.BatchNorm2d(128)\n        self.act3 = nn.ReLU()\n        self.deconv4 = nn.ConvTranspose2d(in_channels = 128, out_channels= 3, kernel_size=(4,4), stride = 1, padding = 1, bias = True)\n        self.bn4= nn.BatchNorm2d(3)\n        self.act4= nn.Sigmoid()\n    def forward(self, x: torch.Tensor):\n        x = self.dense0(x)\n        x = self.act0(x)\n        #print(x.shape)\n        x = self.deconv1(x.view(-1, 16, 8, 8))\n        x = self.bn1(x)\n        x = self.act1(x)\n        #print(x.shape)\n        x = self.deconv2(x)\n        x = self.bn2(x)\n        x = self.act2(x)\n        #print(x.shape)\n        x = self.deconv3(x)\n        x = self.bn3(x)\n        x = self.act3(x)\n        #print(x.shape)\n        x = self.deconv4(x)\n        x = self.bn4(x)\n        x = self.act4(x)\n        #print(x.shape)\n        x = x[:, :, 7:71, 7:71]\n        #print(x.shape)\n        return x","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"LR_g = 0.0002\nLR_d = 0.0002\nmu = 1\ngamma = 0.001\ncriterion = nn.BCELoss()\n\ndiscriminator = Discriminator().cuda()\ngenerator = Generator().cuda()\n\noptimizerD = torch.optim.Adam(discriminator.parameters(), lr=LR_d, betas=(0.5, 0.999))\noptimizerG = torch.optim.Adam(generator.parameters(), lr=LR_g, betas=(0.5, 0.999))\n\nsteps = 500\noutputs = []\nwith tqdm(total = steps) as pbar:\n    for i, real_images in enumerate(data_loader):\n        if i > steps:\n            break\n        real_images = real_images.cuda()\n        for step in range(20):\n            noise = torch.normal(torch.zeros((real_images.shape[0],6)),1.0).cuda()\n            fake_images = generator.forward(noise).detach().cuda()\n            input_images = torch.cat([real_images, fake_images], dim=0).cuda()\n            target = torch.cat([torch.ones(real_images.shape[0], 1), torch.zeros(fake_images.shape[0], 1)], dim=0).cuda()\n\n            # Update discriminator\n            discriminator.zero_grad()\n            loss = criterion(discriminator.forward(input_images), target)\n            loss.backward()\n            optimizerD.step()\n\n            # Update generator\n            generator.zero_grad()\n            generated = generator.forward(noise)\n            dispersion = (torch.cat([(a.flatten()-b.flatten())**16 for a, b in it.product(generated, generated)]).sum())**(1.0/16.0)\n            loss = criterion(discriminator.forward(torch.cat([real_images, generator.forward(noise)], dim=0)), torch.zeros_like(target))\\\n            - mu*torch.sigmoid(gamma*dispersion)\n            loss.backward()\n            optimizerG.step()\n        \n        pbar.set_postfix({'Dispersion':float(dispersion)})\n            \n        # Save images\n        if i % 5 == 0:\n            outputs.append(fake_images.cpu().permute(0,2,3,1).detach().numpy())\n        pbar.update(1)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"for step in outputs:\n    io.imshow_collection(step[0:2])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"for i in range(outputs[-1].shape[0]):\n    if i %3 == 0:\n        io.imshow_collection(outputs[-1][i:i+2])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","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}