{"cells":[{"metadata":{"trusted":true},"cell_type":"code","source":"%matplotlib inline","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import copy\nfrom itertools import chain\nimport io\nimport math\nfrom multiprocessing import cpu_count\nfrom pathlib import Path\nfrom pdb import set_trace\nimport time\nfrom threading import Thread\nfrom xml.etree import ElementTree\nimport zipfile\n\n# General utils\nfrom allennlp.training.learning_rate_schedulers import CosineWithRestarts\nfrom imageio import imread\nfrom joblib import Parallel, delayed\nimport numpy as np\nimport pandas as pd\nimport PIL.Image\nimport matplotlib.pyplot as plt\nfrom tqdm import tqdm_notebook as tqdm\nfrom scipy.stats import truncnorm\n\n# PyTorch\nimport torch\nimport torch.nn as nn\nimport torch.nn.parallel\nimport torch.backends.cudnn as cudnn\nimport torch.optim as optim\nfrom torch.autograd import Variable\nfrom torch.utils.data import Dataset, DataLoader\nimport torchvision.datasets as dset\nimport torchvision.transforms as transforms\nfrom torchvision.utils import make_grid, save_image","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Kernel global variables initialization\n\nWe initialize random generator seed, override `print` function to duplicate its output into kernel logging stream, and setup a \"watchdog\" that tracks how many time we've spent to run the kernel."},{"metadata":{"trusted":true},"cell_type":"code","source":"HOUR = 3600\nTRAINING_TIME = 0.001 * HOUR\nprint(f'Training time: {TRAINING_TIME}')\n\nclass Watchdog:\n    def __init__(self, max_seconds=TRAINING_TIME):\n        self.start = time.time()\n        self.deadline = max_seconds\n    \n    @property\n    def timeout(self):\n        return self.elapsed >= self.deadline\n       \n    @property\n    def elapsed(self):\n        return time.time() - self.start\n\nwd = Watchdog()\n\nSEED = 1\nnp.random.seed(SEED)\ntorch.manual_seed(SEED)\nANNOTS = Path.cwd().parent/'input'/'annotation'/'Annotation'\nIMAGES = Path.cwd().parent/'input'/'all-dogs'/'all-dogs'\n\ntry:\n    # make sure we patch printing function only once\n    patched\nexcept NameError:\n    patched = True\n    __print__ = print\n    def print(message):\n        import os\n        from datetime import datetime\n        log_message = datetime.now().strftime(f'[Kernel][%Y-%m-%d %H:%M:%S] {message}')\n        os.system(f'echo \\\"{log_message}\\\"')\n        __print__(message)\n        \nclass VisualStyle:\n    \"\"\"Convenience wrapper on top of matplotlib config.\"\"\"\n\n    def __init__(self, config, default=None):\n        if default is None:\n            default = plt.rcParams\n        self.default = default.copy()\n        self.config = config\n\n    def replace(self):\n        plt.rcParams = self.config\n\n    def override(self, extra=None):\n        plt.rcParams.update(self.config)\n        if extra is not None:\n            plt.rcParams.update(extra)\n\n    def restore(self):\n        plt.rcParams = self.default\n\n    def __enter__(self):\n        self.override()\n        return self\n\n    def __exit__(self, exc_type, exc_val, exc_tb):\n        self.restore()\n\n\nclass NotebookStyle(VisualStyle):\n    def __init__(self):\n        super().__init__({\n            'figure.figsize': (11, 8),\n            'axes.titlesize': 20,\n            'axes.labelsize': 18,\n            'xtick.labelsize': 14,\n            'ytick.labelsize': 14,\n            'font.size': 16\n        })\n\nNotebookStyle().override()\n        \nprint(f'Annotations: {ANNOTS}')\nprint(f'Images: {IMAGES}')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Functions to prepare the dataset for training\n\nWe need to convert the dataset into format ready for training. For this purpose, we read and crop images, and concatenate them into a single tensor of (B x C x W x H) format suitable for PyTorch."},{"metadata":{"trusted":true},"cell_type":"code","source":"def parse_annotation(path):\n    root = ElementTree.parse(path).getroot()\n    size = [int(root.find(f'size/{leaf}').text) \n            for leaf in ('width', 'height')] \n    bbox = [int(root.find(f'object/bndbox/{leaf}').text) \n            for leaf in ('xmin', 'ymin', 'xmax', 'ymax')]\n    breed = path.parent.name.split('-')[-1]\n    return {'path': str(path), 'name': path.name, \n            'breed': breed, 'size': size, 'bbox': bbox}\n\ndef enrich_with_image_paths(annotations, images_directory):\n    image_files = {x.stem: x for x in images_directory.iterdir()}\n    enriched_data = []\n    for annot in annotations:\n        image_path = image_files.get(annot['name'], None)\n        if image_path is None:\n            print('Warning: image not found for annotation entry: %s.' % annot['path'])\n            continue\n        annot['image'] = str(image_path)\n        enriched_data.append(annot)\n    return enriched_data\n\ndef load_annotations():\n    return enrich_with_image_paths([\n        parse_annotation(path) \n        for directory in ANNOTS.iterdir() \n        for path in directory.iterdir()\n    ], IMAGES)\n\ndef dog(annot):\n    img = imread(annot['image'])\n    xmin, ymin, xmax, ymax = annot['bbox']\n    cropped = img[ymin:ymax, xmin:xmax]    \n    return cropped\n\ndef chunks(seq, chunk_size=10):\n    n = len(seq)\n    n_chunks = n // chunk_size + int((n % chunk_size) != 0)\n    for i in range(n_chunks):\n        yield seq[i*chunk_size:(i+1)*chunk_size]\n        \ndef resize(image, new_size):\n    return np.array(PIL.Image.fromarray(image).resize(new_size))\n\ndef parallel(func, sequence, func_args=None, n_jobs=None):\n    with Parallel(n_jobs=n_jobs or cpu_count()) as p:\n        func_args = func_args or {}\n        results = p(delayed(func)(item, **func_args) for item in sequence)\n    return results\n\ndef load_single_image(annot, size):\n    cropped = dog(annot)\n    resized = resize(cropped, size)\n    return resized\n\ndef load_dogs_images(annots, size=(64, 64)):\n    return np.stack(parallel(load_single_image, annots, func_args={'size': size}))\n\ndef as_pil_list(dataset):\n    return [PIL.Image.fromarray(image, 'RGB') for image in dataset]\n\nclass Normalizer:\n    def __init__(self, method='tanh', params=None):\n        assert method in ('tanh', 'stats')\n        self.method = method\n        self.params = params or {}\n    \n    def transform(self, dataset):\n        dataset = dataset.float()\n        if self.method == 'tanh':\n            return (dataset - 127.5)/127.5\n        if self.method == 'stats':\n            mean = self.params.get('mean', (0.5, 0.5, 0.5))\n            std = self.params.get('std', (0.5, 0.5, 0.5))\n            return functional.normalize(dataset, mean, std)\n    \n    def inv_transform(self, dataset):\n        if self.method == 'tanh':\n            inv_dataset = (dataset + 1)*127.5\n        if self.method == 'stats':\n            mean = params.get('mean', (0.5, 0.5, 0.5))\n            std = params.get('std', (0.5, 0.5, 0.5))\n            mean, std = [torch.as_tensor(\n                x, dtype=torch.float32, device=dataset.device)]\n            dataset.mul_(std[:, None, None]).add_(mean[:, None, None])\n            inv_dataset = dataset\n        return inv_dataset.long()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Reading the data\n\nWe use the functions defined above to read the data and prepare it for training."},{"metadata":{"trusted":true},"cell_type":"code","source":"print('Reading dogs images and annotations.')\nannots = load_annotations()\nprint(f'Total number of examples: {len(annots)}.')\ndogs = load_dogs_images(annots, (128, 128))\nassert len(dogs) == len(annots)\nprint(f'Dogs dataset shape: {dogs.shape}.')\npils = as_pil_list(dogs)\nprint(f'Numbers of PIL images: {len(pils)}')\ndel dogs, annots","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def show_pil(img, *imgs, n_rows=4):\n    imgs = [img] + list(imgs)\n    n_cols = len(imgs) // n_rows\n    f, axes = plt.subplots(n_rows, n_cols)\n    for img, ax in zip(imgs, axes.flat): \n        ax.imshow(img)\n        ax.axis('off')\n    f.subplots_adjust(wspace=0, hspace=0)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# show_pil(*pils[:16])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class PILDataset:\n    def __init__(self, pil_images, transform=None):\n        self.pil_images = pil_images\n        self.tr = transform or (lambda x: x)\n    def __getitem__(self, i):\n        if isinstance(i, int): return self.tr(self.pil_images[i])\n        elif isinstance(i, (list, np.ndarray)): return [self.tr(self.pil_images[ii]) for ii in i]\n        elif isinstance(i, slice): return [self.tr(img) for img in self.pil_images[i]]\n        raise TypeError(f'unknown index type: {type(i)}')\n    def __len__(self):\n        return len(self.pil_images)\n\nclass RandomCropOfFive:\n    def __init__(self, size):\n        self.five_crop = transforms.FiveCrop(size)\n    def __call__(self, x):\n        [idx] = np.random.randint(0, 4, 1)\n        cropped = self.five_crop(x)[idx]\n        return cropped\n    \ndef show_tensor(t, n_rows=4, denorm=False):\n    if denorm: t = (255 * (t + 1)/2)\n    canvas = make_grid(t).numpy().transpose(1, 2, 0).astype(np.uint8)\n    f, ax = plt.subplots(1, 1)\n    ax.imshow(canvas)\n    ax.axis('off')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# show_tensor(torch.stack(dataset[np.random.randint(0, len(dataset), 64)]), denorm=True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def weights_init(m):\n    classname = m.__class__.__name__\n    if classname.find('Conv') != -1:\n        nn.init.normal_(m.weight.data, 0.0, 0.02)\n    elif classname.find('BatchNorm') != -1:\n        nn.init.normal_(m.weight.data, 1.0, 0.02)\n        nn.init.constant_(m.bias.data, 0)\n        \nclass SpectralNorm(nn.Module):\n    def __init__(self, module):\n        super().__init__()\n        self.module = nn.utils.spectral_norm(module)\n    def forward(self, x):\n        return self.module(x)\n\nclass PixelwiseNorm(nn.Module):\n    def __init__(self, alpha=1e-8):\n        super().__init__()\n        self.alpha = alpha\n    def forward(self, x):\n        y = x.pow(2.).mean(dim=1, keepdim=True).add(self.alpha).sqrt()\n        y = x / y\n        return y\n    \nclass MinibatchStdDev(nn.Module):\n    def __init__(self, alpha=1e-8):\n        super().__init__()\n        self.alpha = alpha\n    def forward(self, x):\n        batch_size, _, height, width = x.shape\n        y = x - x.mean(dim=0, keepdim=True)\n        y = y.pow(2.).mean(dim=0, keepdim=False).add(self.alpha).sqrt()\n        y = y.mean().view(1, 1, 1, 1)\n        y = y.repeat(batch_size, 1, height, width)\n        y = torch.cat([x, y], 1)\n        return y\n    \nclass Mixup:\n    def __init__(self, alpha=0.2):\n        self.alpha = alpha\n    def __call__(self, b1, b2): \n        assert b1.size(0) == b2.size(0)\n        lam = np.random.beta(self.alpha, self.alpha, size=b1.size(0))\n        lam = torch.from_numpy(lam).float().to(b1.device)\n        lam = lam.view(-1, 1, 1, 1)\n        return lam*b1 + (1 - lam)*b2","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class DCGAN_D(nn.Module):\n    def __init__(self):\n        super().__init__()\n        self.main = nn.Sequential(\n            SpectralNorm(nn.Conv2d(3, 64, 4, 2, 1, bias=False)),\n            nn.LeakyReLU(0.2, inplace=True),\n            SpectralNorm(nn.Conv2d(64, 128, 4, 2, 1, bias=False)),\n            nn.BatchNorm2d(128),\n            nn.LeakyReLU(0.2, inplace=True),\n            SpectralNorm(nn.Conv2d(128, 256, 4, 2, 1, bias=False)),\n            nn.BatchNorm2d(256),\n            nn.LeakyReLU(0.2, inplace=True),\n            SpectralNorm(nn.Conv2d(256, 512, 4, 2, 1, bias=False)),\n            nn.BatchNorm2d(512),\n            nn.LeakyReLU(0.2, inplace=True),\n            SpectralNorm(nn.Conv2d(512, 1024, 4, 2, 1, bias=False)),\n            nn.BatchNorm2d(1024),\n            nn.LeakyReLU(0.2, inplace=True),\n            MinibatchStdDev(),\n            SpectralNorm(nn.Conv2d(1024 + 1, 1, 4, 2, 1, bias=False)),\n            nn.Sigmoid()\n        )\n        self.main.apply(weights_init)\n    def forward(self, x):\n        return self.main(x)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class DCGAN_G(nn.Module):\n    def __init__(self, nz):\n        super().__init__()\n        self.nz = nz\n        self.main = nn.Sequential(\n            SpectralNorm(nn.ConvTranspose2d(nz, 1024, 4, 1, 0, bias=False)),\n            nn.BatchNorm2d(1024),\n            nn.LeakyReLU(inplace=True),\n            SpectralNorm(nn.ConvTranspose2d(1024, 512, 4, 2, 1, bias=False)),\n            nn.BatchNorm2d(512),\n            nn.LeakyReLU(inplace=True),\n            SpectralNorm(nn.ConvTranspose2d(512, 256, 4, 2, 1, bias=False)),\n            nn.BatchNorm2d(256),\n            nn.LeakyReLU(inplace=True),\n            SpectralNorm(nn.ConvTranspose2d(256, 128, 4, 2, 1, bias=False)),\n            nn.BatchNorm2d(128),\n            nn.LeakyReLU(inplace=True),\n            SpectralNorm(nn.ConvTranspose2d(128, 64, 4, 2, 1, bias=False)),\n            nn.BatchNorm2d(64),\n            nn.LeakyReLU(inplace=True),\n            PixelwiseNorm(),\n            SpectralNorm(nn.ConvTranspose2d(64, 3, 3, 1, 1, bias=False)),\n            nn.Tanh()\n        )\n        self.main.apply(weights_init)\n    def forward(self, x):\n        return self.main(x)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"bs = 16\nnz = 128\nlr_d = 0.0005\nlr_g = 0.0005\nbeta_1 = 0.5\nuse_adam = True\nmixup = Mixup(0.2)\n\ndataset = PILDataset(pils, transform=transforms.Compose([\n    transforms.Resize(70),\n    RandomCropOfFive(64),\n    transforms.RandomHorizontalFlip(),\n    transforms.RandomRotation(3),\n    transforms.ToTensor(),\n    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))\n]))\n\nnetD = DCGAN_D().cuda()\nnetG = DCGAN_G(nz).cuda()\nif use_adam:\n    optimizerD = optim.Adam(netD.parameters(), lr=lr_d, betas=(beta_1, 0.999))\n    optimizerG = optim.Adam(netG.parameters(), lr=lr_g, betas=(beta_1, 0.999))\nelse:\n    optimizerD = optim.RMSprop(netD.parameters(), lr=lr_d)\n    optimizerG = optim.RMSprop(netG.parameters(), lr=lr_g)\nschedD = CosineWithRestarts(optimizerD, eta_min=lr_d*0.1, t_initial=1000, t_mul=math.sqrt(2))\nschedG = CosineWithRestarts(optimizerG, eta_min=lr_g*0.1, t_initial=1000, t_mul=math.sqrt(2))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def truncated_normal(size, threshold=1):\n    return truncnorm.rvs(-threshold, threshold, size=size)\n\ndef sample(dataset, batch_size):\n    idx = np.random.randint(0, len(dataset), batch_size)\n    return torch.stack(dataset[idx]).cuda()\n\ndef smooth_positive(labels):\n    jitter = torch.from_numpy(np.random.uniform(0.05, 0.1, len(labels))).float().to(labels.device)\n    jitter = jitter.view(labels.size())\n    return (labels - jitter)\n\ndef ema(avg_net, curr_net, alpha=0.999):\n    for avg_param, curr_param in zip(avg_net.parameters(), curr_net.parameters()):\n        avg_param.data.mul_(alpha).add_(1 - alpha, curr_param.data)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def train(bs=32):\n    print('Starting training loop...')\n    epoch = 0\n    real_label = 1\n    fake_label = 0\n    loss_fn = nn.BCELoss()\n    n = len(dataset)\n    n_batches = n // bs \n    avg_model = type(netG)(netG.nz).cuda()\n    avg_model.load_state_dict(netG.state_dict())\n    \n    while True:\n        idx1 = np.random.permutation(n)\n        idx2 = np.random.permutation(n)\n        \n        for i in range(n_batches):\n\n            if wd.timeout: return avg_model\n            \n            epoch += 1\n            \n            batch1 = torch.stack(dataset[idx1[i*bs:(i+1)*bs]]).float().cuda()\n            batch2 = torch.stack(dataset[idx2[i*bs:(i+1)*bs]]).float().cuda()\n            mixed = mixup(batch1, batch2)\n            \n            netD.zero_grad()\n            x_real = mixed.cuda()\n            batch_size = x_real.size(0)\n            labels = torch.full((batch_size, 1), real_label).cuda()\n            labels = smooth_positive(labels) \n            output = netD(x_real).view(-1, 1)\n            errD_real = loss_fn(output, labels)\n            errD_real.backward()\n            d_x = output.mean().item()\n\n            noise = torch.from_numpy(truncated_normal((batch_size, nz, 1, 1))).float().cuda()\n            x_fake = netG(noise)\n            labels.fill_(fake_label)\n            output = netD(x_fake.detach()).view(-1, 1)\n            errD_fake = loss_fn(output, labels)\n            errD_fake.backward()\n            d_g_z1 = output.mean().item()\n            errD = errD_real + errD_fake\n            optimizerD.step()\n\n            netG.zero_grad()\n            labels.fill_(real_label)\n            output = netD(x_fake).view(-1, 1)\n            errG = loss_fn(output, labels)\n            errG.backward()\n            d_g_z2 = output.mean().item()\n            optimizerG.step()\n            \n            if epoch % 150 == 0:\n                print(f'[{epoch:06d}][{i+1:03d}] '\n                      f'lr_d: {schedD.get_values()[0]:.6f}, '\n                      f'lr_g: {schedG.get_values()[0]:.6f} | '\n                      f'loss_d: {errD.item():.4f}, '\n                      f'loss_g: {errG.item():.4f} | '\n                      f'D(x): {d_x:.4f}, D(G(z)): {d_g_z1:.4f}/{d_g_z2:.4f}')\n        \n            schedD.step()\n            schedG.step()\n            ema(avg_model, netG)\n            \n    return avg_model","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"avgG = train(bs=16)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print('Final model images generation.')\nprint('Creating archive to write the images.')\narch = zipfile.ZipFile('images.zip', 'w')\nimg_no = 0\nfor batch in range(100):\n    t_noise = torch.from_numpy(truncated_normal((100, nz, 1, 1))).float().cuda()\n    images = avgG(t_noise).detach().cpu()\n    images = images.mul(0.5).add(0.5)\n    images = (255 * images.numpy()).astype(np.uint8)\n    images = images.transpose(0, 2, 3, 1)\n    for image in images:\n        buf = io.BytesIO()\n        PIL.Image.fromarray(image).save(buf, format='png')\n        buf.seek(0)\n        arch.writestr(f'{img_no}.png', buf.getvalue())\n        img_no += 1\narch.close()\nprint('Saving is done!')","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":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.6.8"}},"nbformat":4,"nbformat_minor":1}