{"cells":[{"metadata":{"trusted":true},"cell_type":"code","source":"%matplotlib inline","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from collections import OrderedDict\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\nfrom sklearn.preprocessing import LabelEncoder\n\n# PyTorch\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\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\n\nclass Watchdog:\n    def __init__(self, max_seconds=8.5 * HOUR):\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]","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":{},"cell_type":"markdown","source":"### Layers Utils"},{"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 Reshape(nn.Module):\n    def __init__(self, shape):\n        super().__init__()\n        self.shape = shape\n    def forward(self, x):\n        return x.view(x.size(0), *self.shape)\n    def __repr__(self):\n        return f'{self.__class__.__name__}(shape={self.shape})'\n    \nclass Flatten(nn.Module):\n    def forward(self, x):\n        return x.view(x.size(0), -1)\n\nclass Block(nn.Module):\n    def __init__(\n        self, ni, no,\n        k=3, s=1, p=1, \n        scale=None, avgpool=None,\n        ws=False, gain_ws=2, bn=False, pn=False, \n        activ=nn.LeakyReLU\n    ):\n        super().__init__()\n        module = nn.Conv2d(ni, no, kernel_size=k, stride=s, padding=p, bias=not bn)\n        if ws: module = WeightScaler(module, gain=gain_ws) \n        layers = [module]\n        if scale is not None:\n            layers.insert(0, nn.UpsamplingNearest2d(scale_factor=scale))\n        if bn: layers.append(nn.BatchNorm2d(no))\n        if activ is not None:\n            if activ == nn.ReLU:\n                layers.append(activ(True))\n            elif activ == nn.LeakyReLU:\n                layers.append(activ(negative_slope=0.2, inplace=True))\n            elif activ == nn.PReLU:\n                layers.append(activ(num_parameters=1))\n            else:\n                raise NotImplementedError('unknown activation layer')\n        if pn: layers.append(PixelwiseNorm())\n        if avgpool is not None:\n            layers.append(nn.AvgPool2d(avgpool))\n        self.block = nn.ModuleList(layers)\n        \n    def forward(self, x):\n        for layer in self.block:\n            x = layer(x)\n        return x\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 WeightScaler(nn.Module):\n    def __init__(self, wrapped, gain=2):\n        super().__init__()\n        self.wrapped = wrapped\n        self.gain = gain\n        self.scale = (self.gain / wrapped.weight[0].numel()) ** 0.5\n    def forward(self, x):\n        return self.wrapped(x) * self.scale\n    def __repr__(self):\n        return f'{self.__class__.__name__}({self.wrapped}, c={self.gain})'\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":{},"cell_type":"markdown","source":"### GAN Models"},{"metadata":{"trusted":true},"cell_type":"code","source":"class Generator(nn.Module):\n    def __init__(self, latent_dim):\n        super().__init__()\n        self.to_rgbs = self.create_to_rgb_blocks()\n        self.base = nn.Sequential(\n            PixelwiseNorm(),\n            nn.Linear(latent_dim, 4096),\n            nn.LeakyReLU(0.2, True),\n            PixelwiseNorm(),\n            Reshape((256, 4, 4)),\n            nn.Conv2d(256, 512, 3, 1, 1),\n            nn.LeakyReLU(0.2, True),\n            PixelwiseNorm()\n        )  # output shape: B x 512 x 1 x 1\n        self.conv_1 = nn.ModuleList([\n            Block(512, 256, k=3, s=1, p=1, pn=True, scale=2),\n            Block(256, 128, k=3, s=1, p=1, pn=True, scale=2),\n            Block(128,  64, k=3, s=1, p=1, pn=True, scale=2),\n            Block( 64,  32, k=3, s=1, p=1, pn=True, scale=2)\n        ])\n        self.conv_2 = nn.ModuleList([\n            Block(256, 256, k=3, s=1, p=1, pn=True),\n            Block(128, 128, k=3, s=1, p=1, pn=True),\n            Block( 64,  64, k=3, s=1, p=1, pn=True),\n            Block( 32,  32, k=3, s=1, p=1, pn=True)            \n        ])\n        self.tanh = nn.Tanh()\n        self.apply(weights_init)\n        \n    def create_to_rgb_blocks(self):\n        def to_rgb(n): return nn.Conv2d(n, 3, kernel_size=3, stride=1, padding=1)\n        layers = [to_rgb(x) for x in (512, 256, 128, 64, 32)]\n        return nn.ModuleList(layers)\n    \n    def forward(self, z, block=0, alpha=0.0):\n        x = self.base(z)\n        if block == 0:\n            return self.tanh(self.to_rgbs[0](x))\n        prev = None\n        for i in range(block):\n            prev = x\n            x = self.conv_1[i](x)\n            x = self.conv_2[i](x)\n        # upper branch\n        a = self.to_rgbs[block](x)\n        # lower branch\n        b = F.upsample_nearest(self.to_rgbs[block-1](prev), scale_factor=2)\n        blended = alpha*a + (1-alpha)*b\n        return self.tanh(blended)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class Critic(nn.Module):\n    def __init__(self):\n        super().__init__()\n        self.from_rgbs = self.create_from_rgb_blocks()\n        self.conv_1 = nn.ModuleList([\n            Block( 32,  32, k=3, s=1, p=1),\n            Block( 64,  64, k=3, s=1, p=1),\n            Block(128, 128, k=3, s=1, p=1),\n            Block(256, 256, k=3, s=1, p=1),\n        ])\n        self.conv_2 = nn.ModuleList([\n            Block( 32,  64, k=3, s=1, p=1, avgpool=2),\n            Block( 64, 128, k=3, s=1, p=1, avgpool=2),\n            Block(128, 256, k=3, s=1, p=1, avgpool=2),\n            Block(256, 512, k=3, s=1, p=1, avgpool=2),\n        ])\n        self.top = nn.Sequential(\n            MinibatchStdDev(),\n            nn.Conv2d(513, 512, 3, 1, 1),\n            nn.LeakyReLU(0.2, True),\n            nn.Conv2d(512, 512, 4, 2, 0),\n            nn.LeakyReLU(0.2, True),\n            Flatten(),\n            nn.Linear(512, 1),\n            nn.Sigmoid()\n        )\n    \n    def create_from_rgb_blocks(self):\n        def from_rgb(n): return nn.Conv2d(3, n, kernel_size=3, stride=1, padding=1)\n        layers = [from_rgb(x) for x in (32, 64, 128, 256, 512)]\n        return nn.ModuleList(layers)\n    \n    def forward(self, x, block=0, alpha=0.0):\n        act = lambda value: F.leaky_relu(value, 0.2)\n        \n        if block == 0:\n            # input: 4x4\n            x = F.leaky_relu(self.from_rgbs[-1](x), 0.2)\n            x = self.top(x)\n            return x\n        \n        elif block == 1:\n            # input: 8x8\n            \n            # upper branch\n            a = act(self.from_rgbs[-2](x))\n            a = self.conv_1[-1](a)\n            a = self.conv_2[-1](a)\n            \n            # lower branch\n            b = act(self.from_rgbs[-1](F.avg_pool2d(x, 2)))\n            \n            blended = alpha*a + (1-alpha)*b\n            x = self.top(blended)\n            return x\n        \n        elif block == 2:\n            # input: 16x16\n            \n            # upper branch\n            a = act(self.from_rgbs[-3](x))\n            a = self.conv_1[-2](a)\n            a = self.conv_2[-2](a)\n            \n            # lower branch\n            b = act(self.from_rgbs[-2](F.avg_pool2d(x, 2)))\n            \n            blended = alpha*a + (1-alpha)*b\n            \n            x = self.conv_1[-1](blended)\n            x = self.conv_2[-1](x)\n            x = self.top(x)\n            return x\n        \n        elif block == 3:\n            # input: 32x32\n            \n            # upper branch\n            a = act(self.from_rgbs[-4](x))\n            a = self.conv_1[-3](a)\n            a = self.conv_2[-3](a)\n            \n            # lower branch\n            b = act(self.from_rgbs[-3](F.avg_pool2d(x, 2)))\n            \n            blended = alpha*a + (1-alpha)*b\n            \n            x = self.conv_1[-2](blended)\n            x = self.conv_2[-2](x)\n            x = self.conv_1[-1](x)\n            x = self.conv_2[-1](x)\n            x = self.top(x)\n            return x\n        \n        elif block == 4:\n            # input: 64x64\n            \n            # upper branch\n            a = act(self.from_rgbs[-5](x))\n            a = self.conv_1[-4](a)\n            a = self.conv_2[-4](a)\n            \n            # lower branch\n            b = act(self.from_rgbs[-4](F.avg_pool2d(x, 2)))\n            \n            blended = alpha*a + (1-alpha)*b\n            \n            x = self.conv_1[-3](blended)\n            x = self.conv_2[-3](x)\n            x = self.conv_1[-2](x)\n            x = self.conv_2[-2](x)\n            x = self.conv_1[-1](x)\n            x = self.conv_2[-1](x)\n            x = self.top(x)\n            return x","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# d = Critic()\n# img = torch.ones((10, 3, 4, 4))\n# out = d(img, block=0, alpha=0.5)\n# out.shape","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# d = Critic()\n# img = torch.ones((10, 3, 8, 8))\n# out = d(img, block=1, alpha=0.5)\n# out.shape","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# d = Critic()\n# img = torch.ones((10, 3, 16, 16))\n# out = d(img, block=2, alpha=0.5)\n# out.shape","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# d = Critic()\n# img = torch.ones((10, 3, 32, 32))\n# out = d(img, block=3, alpha=0.5)\n# out.shape","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# d = Critic()\n# img = torch.ones((10, 3, 64, 64))\n# out = d(img, block=4, alpha=0.5)\n# out.shape","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# g = Generator(128)\n# z = torch.randn((10, 128))\n# x = g(z)\n# x.shape","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Training"},{"metadata":{"trusted":true},"cell_type":"code","source":"bs = 16\nnz = 128\nlr_d = 0.0001\nlr_g = 0.0001\nbeta_1 = 0.5\nbeta_2 = 0.999\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 = Critic().cuda()\noptimizerD = optim.Adam(netD.parameters(), lr=lr_d, betas=(beta_1, beta_2))\nschedD = CosineWithRestarts(optimizerD, eta_min=lr_d*0.1, t_initial=1000, t_mul=math.sqrt(2))\n                        \nnetG = Generator(nz).cuda()\noptimizerG = optim.Adam(netG.parameters(), lr=lr_g, betas=(beta_1, beta_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":"# 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 truncated_normal(size, threshold=1):\n    return truncnorm.rvs(-threshold, threshold, size=size)\n\ndef hypersphere(z, radius=1):\n    return z * radius / z.norm(p=2, dim=1, keepdim=True)\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)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# def train(bs=32):\ndef train():\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    block = 0\n    alpha_range = list(np.linspace(0, 1, 7500)) + list([1] * 7500)\n    alpha_index = 0\n    # batch_sizes = {0: 256, 1: 128, 2: 64, 3: 32, 4: 16}\n    batch_sizes = {0: 512, 1: 256, 2: 128, 3: 32, 4: 16}\n    \n    while True:\n        bs = batch_sizes[block]\n        n_batches = n // bs\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\n            \n            epoch += 1\n            alpha_value = alpha_range[alpha_index]\n            \n            batch1 = torch.stack(dataset[idx1[i*bs:(i+1)*bs]]).float()\n            batch2 = torch.stack(dataset[idx2[i*bs:(i+1)*bs]]).float()\n            mixed = mixup(batch1, batch2)\n            \n            if block < 4:\n                mixed = nn.functional.interpolate(mixed, size=(\n                    (4, 4) if block == 0 else\n                    (8, 8) if block == 1 else\n                    (16, 16) if block == 2 else\n                    (32, 32) if block == 3 else\n                    (64, 64)\n                ), mode='bilinear', align_corners=True).detach()\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, block=block, alpha=alpha_value).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            noise = torch.from_numpy(truncated_normal((batch_size, nz))).float().cuda()\n            x_fake = netG(noise, block=block, alpha=alpha_value)\n            labels.fill_(fake_label)\n            output = netD(x_fake.detach(), block=block, alpha=alpha_value).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, block=block, alpha=alpha_value).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:6d}] '\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                      f'no={block}, a={alpha_value:.4f}, bs={bs:d}, sz={list(mixed.shape[2:])}')\n        \n            if block == 4:\n                schedD.step()\n                schedG.step()\n            \n            alpha_index += 1\n            \n            if block == 4:\n                # discard blending, use layers only\n                alpha_index = -1\n            elif alpha_index >= len(alpha_range):\n                # start new cycle\n                block += 1 \n                alpha_index = 0","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# train(bs=16)\ntrain()","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    t_noise = torch.from_numpy(truncated_normal((100, nz))).float().cuda()\n    images = netG(t_noise, block=4, alpha=1.0).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":{"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.7.3"},"toc":{"base_numbering":1,"nav_menu":{},"number_sections":true,"sideBar":true,"skip_h1_title":false,"title_cell":"Table of Contents","title_sidebar":"Contents","toc_cell":false,"toc_position":{},"toc_section_display":true,"toc_window_display":false}},"nbformat":4,"nbformat_minor":1}