{"cells":[{"metadata":{"_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5"},"cell_type":"markdown","source":"# GAN mix\n\nCombining ideas from:\n- [dogs-starter-24jul](https://www.kaggle.com/phoenix9032/gan-dogs-starter-24-jul-custom-layers)\n- [Pytorch RaLS-C-SAGAN](https://www.kaggle.com/mpalermo/pytorch-rals-c-sagan)"},{"metadata":{"trusted":true},"cell_type":"code","source":"import gzip\nimport os\nimport pathlib\nimport pickle\nimport random\nimport shutil\nimport time\nimport urllib\nimport warnings\nimport xml.etree.ElementTree as ET\nimport zipfile\nfrom time import time\n\nimport imgaug as ia\nimport imgaug.augmenters as iaa\nimport matplotlib.image as mpimg\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport pandas as pd\nimport PIL\nimport tensorflow as tf\nimport torch as th\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.nn.parallel\nimport torch.optim as optim\nimport torch.utils.data\nimport torchvision\nimport torchvision.datasets as dset\nimport torchvision.transforms as transforms\nimport torchvision.utils as vutils\nfrom PIL import Image\nfrom scipy import linalg\nfrom scipy.stats import truncnorm\nfrom torch import nn as nn\nfrom torch import optim as optim\nfrom torch.autograd import Variable\nfrom torch.nn import Parameter\nfrom torch.nn.init import xavier_uniform_\nfrom torch.nn.utils import spectral_norm\nfrom torch.utils.data import DataLoader, Dataset\nfrom torchvision import datasets as dset\nfrom torchvision import transforms as transforms\nfrom torchvision.utils import save_image\nfrom tqdm import tqdm\n\nprint(os.listdir(\"../input\"))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# utilities\ndef seed_everything(seed=42):\n    random.seed(seed)\n    os.environ[\"PYTHONHASHSEED\"] = str(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.backends.cudnn.deterministic = True\n    return\n\n\ndef show_generated_img_all():\n    gen_z = torch.randn(32, nz, 1, 1, device=device)\n    gen_images = netG(gen_z).to(\"cpu\").clone().detach()\n    gen_images = gen_images.numpy().transpose(0, 2, 3, 1)\n    gen_images = (gen_images + 1.0) / 2.0\n    fig = plt.figure(figsize=(25, 16))\n    for ii, img in enumerate(gen_images):\n        ax = fig.add_subplot(4, 8, ii + 1, xticks=[], yticks=[])\n        plt.imshow(img)\n    # plt.savefig(filename)\n\n\n### This is to show one sample image for iteration of chosing\ndef 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    gen_image = (gen_image + 1.0) / 2.0\n    plt.imshow(gen_image)\n    plt.show()\n\n\ndef mse(imageA, imageB):\n    err = np.sum((imageA.astype(\"float\") - imageB.astype(\"float\")) ** 2)\n    err /= float(imageA.shape[0] * imageA.shape[1])\n    return err\n\n\ndef show_generated_img(n_images=5, nz=128):\n    sample = []\n    for _ in range(n_images):\n        noise = torch.randn(1, nz, 1, 1, device=device)\n        dog_label = torch.randint(0, len(encoded_dog_labels), (1,), device=device)\n        gen_image = netG((noise, dog_label)).to(\"cpu\").clone().detach().squeeze(0)\n        gen_image = gen_image.numpy().transpose(1, 2, 0)\n        sample.append(gen_image)\n\n    figure, axes = plt.subplots(1, len(sample), figsize=(64, 64))\n    for index, axis in enumerate(axes):\n        axis.axis(\"off\")\n        image_array = (sample[index] + 1.0) / 2.0\n        axis.imshow(image_array)\n    plt.show()\n\n\ndef analyse_generated_by_class(n_images=5):\n    good_breeds = []\n    for l in range(len(decoded_dog_labels)):\n        sample = []\n        for _ in range(n_images):\n            noise = torch.randn(1, nz, 1, 1, device=device)\n            dog_label = torch.full((1,), l, device=device, dtype=torch.long)\n            gen_image = netG((noise, dog_label)).to(\"cpu\").clone().detach().squeeze(0)\n            gen_image = gen_image.numpy().transpose(1, 2, 0)\n            sample.append(gen_image)\n\n        d = np.round(\n            np.sum([mse(sample[k], sample[k + 1]) for k in range(len(sample) - 1)])\n            / n_images,\n            1,\n        )\n        if d < 1.0:\n            continue  # had mode colapse(discard)\n\n        print(f\"Generated breed({d}): \", decoded_dog_labels[l])\n        figure, axes = plt.subplots(1, len(sample), figsize=(64, 64))\n        for index, axis in enumerate(axes):\n            axis.axis(\"off\")\n            image_array = (sample[index] + 1.0) / 2.0\n            axis.imshow(image_array)\n        plt.show()\n\n        good_breeds.append(l)\n    return good_breeds\n\n\ndef create_submit(good_breeds):\n    print(\"Creating submit\")\n    os.makedirs(\"../output_images\", exist_ok=True)\n    im_batch_size = 100\n    n_images = 10000\n\n    all_dog_labels = np.random.choice(good_breeds, size=n_images, replace=True)\n    for i_batch in range(0, n_images, im_batch_size):\n        noise = torch.randn(im_batch_size, nz, 1, 1, device=device)\n        dog_labels = torch.from_numpy(\n            all_dog_labels[i_batch : (i_batch + im_batch_size)]\n        ).to(device)\n        gen_images = netG((noise, dog_labels))\n        gen_images = (gen_images.to(\"cpu\").clone().detach() + 1) / 2\n        for ii, img in enumerate(gen_images):\n            save_image(\n                gen_images[ii, :, :, :],\n                os.path.join(\"../output_images\", f\"image_{i_batch + ii:05d}.png\"),\n            )\n\n    import shutil\n\n    shutil.make_archive(\"images\", \"zip\", \"../output_images\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class DataGenerator(Dataset):\n    def __init__(self, directory, transform=None, n_samples=np.inf, crop_dogs=True):\n        self.directory = directory\n        self.transform = transform\n        self.n_samples = n_samples\n        self.samples, self.labels = self.load_dogs_data(directory, crop_dogs)\n\n    def load_dogs_data(self, directory, crop_dogs):\n        required_transforms = torchvision.transforms.Compose(\n            [torchvision.transforms.Resize(64), torchvision.transforms.CenterCrop(64)]\n        )\n\n        imgs = []\n        labels = []\n        paths = []\n        for root, _, fnames in sorted(os.walk(directory)):\n            for fname in sorted(fnames)[: min(self.n_samples, 1e7)]:\n                path = os.path.join(root, fname)\n                paths.append(path)\n        if LOCAL:\n            ROOT = \"../input/Annotation/Annotation/\"\n        else:\n            ROOT = \"../input/annotation/Annotation/\"\n\n        for path in paths:\n            # Load image\n            try:\n                img = dset.folder.default_loader(path)\n            except:\n                continue\n\n            # Get bounding boxes\n            annotation_basename = os.path.splitext(os.path.basename(path))[0]\n            annotation_dirname = next(\n                dirname\n                for dirname in os.listdir(ROOT)\n                if dirname.startswith(annotation_basename.split(\"_\")[0])\n            )\n\n            if crop_dogs:\n                tree = ET.parse(\n                    os.path.join(\n                        ROOT,\n                        annotation_dirname,\n                        annotation_basename,\n                    )\n                )\n                root = tree.getroot()\n                objects = root.findall(\"object\")\n                for o in objects:\n                    bndbox = o.find(\"bndbox\")\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                    object_img = required_transforms(img.crop((xmin, ymin, xmax, ymax)))\n                    imgs.append(object_img)\n                    labels.append(annotation_dirname.split(\"-\")[1].lower())\n\n            else:\n                object_img = required_transforms(img)\n                imgs.append(object_img)\n                labels.append(annotation_dirname.split(\"-\")[1].lower())\n\n        return imgs, labels\n\n    def __getitem__(self, index):\n        sample = self.samples[index]\n        label = self.labels[index]\n\n        if self.transform is not None:\n            sample = self.transform(sample)\n        return np.asarray(sample), label\n\n    def __len__(self):\n        return len(self.samples)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# model utilities\n# ----------------------------------------------------------------------------\n# Pixelwise feature vector normalization.\n# reference: https://github.com/tkarras/progressive_growing_of_gans/blob/master/networks.py#L120\n# ----------------------------------------------------------------------------\nclass PixelwiseNorm(nn.Module):\n    def __init__(self):\n        super(PixelwiseNorm, self).__init__()\n\n    def forward(self, x, alpha=1e-8):\n        \"\"\"\n        forward pass of the module\n        :param x: input activations volume\n        :param alpha: small number for numerical stability\n        :return: y => pixel normalized activations\n        \"\"\"\n        y = x.pow(2.0).mean(dim=1, keepdim=True).add(alpha).sqrt()  # [N1HW]\n        y = x / y  # normalize the input x volume\n        return y\n\n\nclass MinibatchStdDev(th.nn.Module):\n    \"\"\"\n    Minibatch standard deviation layer for the discriminator\n    \"\"\"\n\n    def __init__(self):\n        \"\"\"\n        derived class constructor\n        \"\"\"\n        super(MinibatchStdDev, self).__init__()\n\n    def forward(self, x, alpha=1e-8):\n        \"\"\"\n        forward pass of the layer\n        :param x: input activation volume\n        :param alpha: small number for numerical stability\n        :return: y => x appended with standard deviation constant map\n        \"\"\"\n        batch_size, _, height, width = x.shape\n        # [B x C x H x W] Subtract mean over batch.\n        y = x - x.mean(dim=0, keepdim=True)\n        # [1 x C x H x W]  Calc standard deviation over batch\n        y = th.sqrt(y.pow(2.0).mean(dim=0, keepdim=False) + alpha)\n\n        # [1]  Take average over feature_maps and pixels.\n        y = y.mean().view(1, 1, 1, 1)\n\n        # [B x 1 x H x W]  Replicate over group and pixels.\n        y = y.repeat(batch_size, 1, height, width)\n\n        # [B x C x H x W]  Append as new feature_map.\n        y = th.cat([x, y], 1)\n        # return the computed values:\n        return y\n\n\ndef snconv2d(\n    in_channels,\n    out_channels,\n    kernel_size,\n    stride=1,\n    padding=0,\n    dilation=1,\n    groups=1,\n    bias=True,\n):\n    return spectral_norm(\n        nn.Conv2d(\n            in_channels=in_channels,\n            out_channels=out_channels,\n            kernel_size=kernel_size,\n            stride=stride,\n            padding=padding,\n            dilation=dilation,\n            groups=groups,\n            bias=bias,\n        )\n    )\n\n\ndef snlinear(in_features, out_features):\n    return spectral_norm(nn.Linear(in_features=in_features, out_features=out_features))\n\n\ndef sn_embedding(num_embeddings, embedding_dim):\n    return spectral_norm(\n        nn.Embedding(num_embeddings=num_embeddings, embedding_dim=embedding_dim)\n    )\n\n\nclass Self_Attn(nn.Module):\n    \"\"\" Self attention Layer\"\"\"\n\n    def __init__(self, in_channels):\n        super(Self_Attn, self).__init__()\n        self.in_channels = in_channels\n        self.snconv1x1_theta = snconv2d(\n            in_channels=in_channels,\n            out_channels=in_channels // 8,\n            kernel_size=1,\n            stride=1,\n            padding=0,\n        )\n        self.snconv1x1_phi = snconv2d(\n            in_channels=in_channels,\n            out_channels=in_channels // 8,\n            kernel_size=1,\n            stride=1,\n            padding=0,\n        )\n        self.snconv1x1_g = snconv2d(\n            in_channels=in_channels,\n            out_channels=in_channels // 2,\n            kernel_size=1,\n            stride=1,\n            padding=0,\n        )\n        self.snconv1x1_attn = snconv2d(\n            in_channels=in_channels // 2,\n            out_channels=in_channels,\n            kernel_size=1,\n            stride=1,\n            padding=0,\n        )\n        self.maxpool = nn.MaxPool2d(2, stride=2, padding=0)\n        self.softmax = nn.Softmax(dim=-1)\n        self.sigma = nn.Parameter(torch.zeros(1))\n\n    def forward(self, x):\n        _, ch, h, w = x.size()\n        # Theta path\n        theta = self.snconv1x1_theta(x)\n        theta = theta.view(-1, ch // 8, h * w)\n        # Phi path\n        phi = self.snconv1x1_phi(x)\n        phi = self.maxpool(phi)\n        phi = phi.view(-1, ch // 8, h * w // 4)\n        # Attn map\n        attn = torch.bmm(theta.permute(0, 2, 1), phi)\n        attn = self.softmax(attn)\n        # g path\n        g = self.snconv1x1_g(x)\n        g = self.maxpool(g)\n        g = g.view(-1, ch // 2, h * w // 4)\n        # Attn_g\n        attn_g = torch.bmm(g, attn.permute(0, 2, 1))\n        attn_g = attn_g.view(-1, ch // 2, h, w)\n        attn_g = self.snconv1x1_attn(attn_g)\n        # Out\n        out = x + self.sigma * attn_g\n        return out\n\n\nclass ConditionalBatchNorm2d(nn.Module):\n    def __init__(self, num_features, num_classes):\n        super().__init__()\n        self.num_features = num_features\n        self.bn = nn.BatchNorm2d(num_features)\n        self.embed = nn.Embedding(num_classes, num_features * 2)\n        self.embed.weight.data[:, :num_features].fill_(1.0)  # Initialize scale to 1\n        self.embed.weight.data[:, num_features:].zero_()  # Initialize bias at 0\n\n    def forward(self, inputs):\n        x, y = inputs\n\n        out = self.bn(x)\n        gamma, beta = self.embed(y).chunk(2, 1)\n        out = gamma.view(-1, self.num_features, 1, 1) * out + beta.view(\n            -1, self.num_features, 1, 1\n        )\n        return out\n\n\ndef truncated_normal(size, threshold=1):\n    values = truncnorm.rvs(-threshold, threshold, size=size)\n    return values","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class UpConvBlock(nn.Module):\n    def __init__(\n        self,\n        n_input,\n        n_output,\n        num_classes,\n        k_size=4,\n        stride=2,\n        padding=0,\n        bias=False,\n        dropout_p=0.0,\n        norm=None,\n    ):\n        super(UpConvBlock, self).__init__()\n        self.norm = norm\n        self.dropout_p = dropout_p\n        self.upconv = spectral_norm(\n            nn.ConvTranspose2d(\n                n_input,\n                n_output,\n                kernel_size=k_size,\n                stride=stride,\n                padding=padding,\n                bias=bias,\n            )\n        )\n        if norm == \"cbn\":\n            self.norm = ConditionalBatchNorm2d(n_output, num_classes)\n        elif norm == \"pixnorm\":\n            self.norm = PixelwiseNorm()\n        elif norm == \"bn\":\n            self.norm = nn.BatchNorm2d(n_output)\n        elif norm == None:\n            self.norm = None\n        self.activ = nn.LeakyReLU(0.05, inplace=True)\n        self.dropout = nn.Dropout2d(p=dropout_p)\n\n    def forward(self, inputs):\n        x0, labels = inputs\n\n        x = self.upconv(x0)\n        if self.norm is not None:\n            if self.norm == \"cbn\":\n                x = self.activ(self.norm((x, labels)))\n            else:\n                x = self.activ(self.norm(x))\n        if self.dropout_p > 0.0:\n            x = self.dropout(x)\n        return x\n\n\nclass Generator(nn.Module):\n    def __init__(self, nz=128, num_classes=120, channels=3, nfilt=64):\n        super(Generator, self).__init__()\n        self.nz = nz\n        self.num_classes = num_classes\n        self.channels = channels\n\n        self.label_emb = nn.Embedding(num_classes, nz)\n        self.pixnorm = PixelwiseNorm()\n\n        self.upconv1 = UpConvBlock(\n            2 * nz,\n            nfilt * 16,\n            num_classes,\n            k_size=4,\n            stride=1,\n            padding=0,\n            dropout_p=0.1,\n        )\n        self.upconv2 = UpConvBlock(\n            nfilt * 16,\n            nfilt * 8,\n            num_classes,\n            k_size=4,\n            stride=2,\n            padding=1,\n            dropout_p=0.1,\n            norm=\"pixnorm\",\n        )\n        self.upconv3 = UpConvBlock(\n            nfilt * 8,\n            nfilt * 4,\n            num_classes,\n            k_size=4,\n            stride=2,\n            padding=1,\n            dropout_p=0.1,\n            norm=\"pixnorm\",\n        )\n        self.upconv4 = UpConvBlock(\n            nfilt * 4,\n            nfilt * 2,\n            num_classes,\n            k_size=4,\n            stride=2,\n            padding=1,\n            dropout_p=0.1,\n            norm=\"pixnorm\",\n        )\n        self.upconv5 = UpConvBlock(\n            nfilt * 2,\n            nfilt,\n            num_classes,\n            k_size=4,\n            stride=2,\n            padding=1,\n            dropout_p=0.1,\n            norm=\"pixnorm\",\n        )\n        self.self_attn = Self_Attn(nfilt)\n        self.upconv6 = UpConvBlock(nfilt, 3, num_classes, k_size=3, stride=1, padding=1)\n        self.out_conv = spectral_norm(nn.Conv2d(3, 3, 3, 1, 1, bias=False))\n        self.out_activ = nn.Tanh()\n\n    def forward(self, inputs):\n        z, labels = inputs\n\n        enc = self.label_emb(labels).view((-1, self.nz, 1, 1))\n        enc = F.normalize(enc, p=2, dim=1)\n        x = torch.cat((z, enc), 1)\n\n        x = self.upconv1((x, labels))\n        x = self.upconv2((x, labels))\n        x = self.upconv3((x, labels))\n        x = self.upconv4((x, labels))\n        x = self.upconv5((x, labels))\n        x = self.self_attn(x)\n        x = self.upconv6((x, labels))\n        x = self.out_conv(x)\n        img = self.out_activ(x)\n        \n        return img\n\n\nclass Discriminator(nn.Module):\n    def __init__(self, num_classes=120, channels=3, nfilt=64, emb_dim=64):\n        super(Discriminator, self).__init__()\n        self.channels = channels\n        self.num_classes = num_classes\n        self.emb_dim = emb_dim\n\n        def down_convlayer(\n            n_input, n_output, k_size=4, stride=2, padding=0, dropout_p=0.0\n        ):\n            block = [\n                spectral_norm(\n                    nn.Conv2d(\n                        n_input,\n                        n_output,\n                        kernel_size=k_size,\n                        stride=stride,\n                        padding=padding,\n                        bias=False,\n                    )\n                ),\n                nn.BatchNorm2d(n_output),\n                nn.LeakyReLU(0.2, inplace=True),\n            ]\n            if dropout_p > 0.0:\n                block.append(nn.Dropout(p=dropout_p))\n            return block\n\n        self.label_emb = nn.Embedding(num_classes, self.emb_dim * self.emb_dim)\n        self.model = nn.Sequential(\n            *down_convlayer(self.channels + 1, nfilt, 4, 2, 1),\n            Self_Attn(nfilt),\n            *down_convlayer(nfilt, nfilt * 2, 4, 2, 1, dropout_p=0.10),\n            *down_convlayer(nfilt * 2, nfilt * 4, 4, 2, 1, dropout_p=0.15),\n            *down_convlayer(nfilt * 4, nfilt * 8, 4, 2, 1, dropout_p=0.25),\n            spectral_norm(nn.Conv2d(nfilt * 8, 1, 4, 1, 0, bias=False)),\n        )\n\n    def forward(self, inputs):\n        imgs, labels = inputs\n\n        enc = self.label_emb(labels).view((-1, 1, self.emb_dim, self.emb_dim))\n        enc = F.normalize(enc, p=2, dim=1)\n        x = torch.cat((imgs, enc), 1)\n        out = self.model(x)\n        out = torch.sigmoid(out)\n        \n        return out.view(-1)\n\n\ndef 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)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"raw","source":"# model definitions\nclass Generator(nn.Module):\n    def __init__(self, nz=128, num_classes=120, nfeats=64, nchannels=3):\n        super(Generator, self).__init__()\n\n        self.label_emb = nn.Embedding(num_classes, nz)\n        # input is Z, going into a convolution\n        self.conv1 = spectral_norm(\n            nn.ConvTranspose2d(nz, nfeats * 8, 4, 1, 0, bias=False)\n        )\n        # self.bn1 = nn.BatchNorm2d(nfeats * 8)\n        # state size. (nfeats*8) x 4 x 4\n\n        self.conv2 = spectral_norm(\n            nn.ConvTranspose2d(nfeats * 8, nfeats * 8, 4, 2, 1, bias=False)\n        )\n        # self.bn2 = nn.BatchNorm2d(nfeats * 8)\n        # state size. (nfeats*8) x 8 x 8\n\n        self.conv3 = spectral_norm(\n            nn.ConvTranspose2d(nfeats * 8, nfeats * 4, 4, 2, 1, bias=False)\n        )\n        # self.bn3 = nn.BatchNorm2d(nfeats * 4)\n        # state size. (nfeats*4) x 16 x 16\n\n        self.conv4 = spectral_norm(\n            nn.ConvTranspose2d(nfeats * 4, nfeats * 2, 4, 2, 1, bias=False)\n        )\n        # self.bn4 = nn.BatchNorm2d(nfeats * 2)\n        # state size. (nfeats * 2) x 32 x 32\n\n        self.conv5 = spectral_norm(\n            nn.ConvTranspose2d(nfeats * 2, nfeats, 4, 2, 1, bias=False)\n        )\n        # self.bn5 = nn.BatchNorm2d(nfeats)\n        # state size. (nfeats) x 64 x 64\n\n        self.conv6 = spectral_norm(\n            nn.ConvTranspose2d(nfeats, nchannels, 3, 1, 1, bias=False)\n        )\n        # state size. (nchannels) x 64 x 64\n        self.pixnorm = PixelwiseNorm()\n\n    def forward(self, inputs):\n\n        z, labels = inputs\n\n        enc = self.label_emb(labels).view((-1, self.nz, 1, 1))\n        enc = F.normalize(enc, p=2, dim=1)\n        x = torch.cat((z, enc), 1)\n        \n        x = F.leaky_relu(self.conv1(x))\n        x = F.leaky_relu(self.conv2(x))\n        x = self.pixnorm(x)\n        x = F.leaky_relu(self.conv3(x))\n        x = self.pixnorm(x)\n        x = F.leaky_relu(self.conv4(x))\n        x = self.pixnorm(x)\n        x = F.leaky_relu(self.conv5(x))\n        x = self.pixnorm(x)\n        x = torch.tanh(self.conv6(x))\n\n        return x\n\n\nclass Discriminator(nn.Module):\n    def __init__(self, nchannels=3, nfeats=64):\n        super(Discriminator, self).__init__()\n\n        # input is (nchannels) x 64 x 64\n        self.conv1 = nn.Conv2d(nchannels, nfeats, 4, 2, 1, bias=False)\n        # state size. (nfeats) x 32 x 32\n\n        self.conv2 = spectral_norm(nn.Conv2d(nfeats, nfeats * 2, 4, 2, 1, bias=False))\n        self.bn2 = nn.BatchNorm2d(nfeats * 2)\n        # state size. (nfeats*2) x 16 x 16\n\n        self.conv3 = spectral_norm(\n            nn.Conv2d(nfeats * 2, nfeats * 4, 4, 2, 1, bias=False)\n        )\n        self.bn3 = nn.BatchNorm2d(nfeats * 4)\n        # state size. (nfeats*4) x 8 x 8\n\n        self.conv4 = spectral_norm(\n            nn.Conv2d(nfeats * 4, nfeats * 8, 4, 2, 1, bias=False)\n        )\n        self.bn4 = nn.MaxPool2d(2)\n        # state size. (nfeats*8) x 4 x 4\n        self.batch_discriminator = MinibatchStdDev()\n        self.pixnorm = PixelwiseNorm()\n        self.conv5 = spectral_norm(nn.Conv2d(nfeats * 8 + 1, 1, 2, 1, 0, bias=False))\n        # state size. 1 x 1 x 1\n\n    def forward(self, x):\n        x = F.leaky_relu(self.conv1(x), 0.2)\n        x = F.leaky_relu(self.bn2(self.conv2(x)), 0.2)\n        # x = self.pixnorm(x)\n        x = F.leaky_relu(self.bn3(self.conv3(x)), 0.2)\n        # x = self.pixnorm(x)\n        x = F.leaky_relu(self.bn4(self.conv4(x)), 0.2)\n        # x = self.pixnorm(x)\n        x = self.batch_discriminator(x)\n        x = torch.sigmoid(self.conv5(x))\n        # x= self.conv5(x)\n        return x.view(-1, 1)"},{"metadata":{"trusted":true},"cell_type":"code","source":"LOCAL = False\nBATCH_SIZE = 32\n\nstart = time()\nseed_everything()\n\n# dataset initialization\nif LOCAL:\n    database = \"../input/all-dogs/\"\nelse:\n    database = \"../input/all-dogs/all-dogs/\"\n\n    \ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\ntransform = transforms.Compose(\n    [\n        transforms.RandomHorizontalFlip(p=0.3),\n        transforms.ToTensor(),\n        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),\n    ]\n)\n\ntrain_data = DataGenerator(database, transform=transform, n_samples=25000)\ntrain_loader = torch.utils.data.DataLoader(\n    train_data, shuffle=True, batch_size=BATCH_SIZE, num_workers=4\n)\n\ndecoded_dog_labels = {\n    i: breed for i, breed in enumerate(sorted(set(train_data.labels)))\n}\nencoded_dog_labels = {\n    breed: i for i, breed in enumerate(sorted(set(train_data.labels)))\n}\ntrain_data.labels = [encoded_dog_labels[l] for l in train_data.labels]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# inspect training samples & labels, first batch\nfor i, (x, y) in enumerate(train_loader):\n    imgs_, labels_ = x, y\n    break\nimgs_ = np.swapaxes(imgs_.numpy(), 1, -1)\nlabels_ = labels_.numpy()\n\nN_COLS, N_ROWS = 4, 4\nfig, ax = plt.subplots(N_COLS, N_ROWS, figsize=(20, 20))\nidx = 0\nfor i in range(N_COLS):\n    for j in range(N_ROWS):\n        ax[i, j].imshow(imgs_[idx])\n        ax[i, j].set_title(labels_[idx])\n        idx += 1","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# training parameters initialization\nLR_DISC = 0.0003\nLR_GEN = 0.0001\nNOISE_DIM = 128\n\nbeta1 = 0.5\nepochs = 301\nreal_label = 0.7\nfake_label = 0.0\n\n\ncriterion = nn.BCELoss()\n# criterion = nn.MSELoss()\n\nnetG = Generator(nz=NOISE_DIM, num_classes=120).to(device)\nnetD = Discriminator(120, 3, 64).to(device)\n\noptimizerD = optim.Adam(netD.parameters(), lr=LR_DISC, betas=(beta1, 0.99))\noptimizerG = optim.Adam(netG.parameters(), lr=LR_GEN, betas=(beta1, 0.99))\n\nlr_schedulerG = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(\n    optimizerG, T_0=epochs // 200, eta_min=0.00005\n)\nlr_schedulerD = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(\n    optimizerD, T_0=epochs // 200, eta_min=0.00005\n)\n\nnz = NOISE_DIM\nfixed_noise = torch.randn(BATCH_SIZE, nz, 1, 1, device=device)\nbatch_size = train_loader.batch_size","execution_count":null,"outputs":[]},{"metadata":{"_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","_kg_hide-input":true,"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","trusted":true},"cell_type":"code","source":"### training here\nstep = 0\nfor epoch in range(epochs):\n    for ii, (real_images, class_labels) in tqdm(enumerate(train_loader), total=len(train_loader)):\n        end = time()\n        if (end - start) > 25000:\n            break\n        ############################\n        # (1) Update D network: maximize log(D(x)) + log(1 - D(G(z)))\n        ###########################\n        # train with real\n        netD.zero_grad()\n        real_images = real_images.to(device)\n        class_labels = torch.tensor(class_labels, device=device)\n        batch_size = real_images.size(0)\n        labels = torch.full(\n            (batch_size, 1), real_label, device=device\n        ) + np.random.uniform(-0.1, 0.1)\n\n        output = netD((real_images, class_labels))\n        errD_real = criterion(output, labels)\n        errD_real.backward()\n        D_x = output.mean().item()\n\n        # train with fake\n        noise = torch.randn(batch_size, nz, 1, 1, device=device)\n        fake = netG((noise, class_labels))\n        labels.fill_(fake_label) + np.random.uniform(0, 0.2)\n        output = netD((fake.detach(), class_labels))\n        errD_fake = criterion(output, labels)\n        errD_fake.backward()\n        D_G_z1 = output.mean().item()\n        errD = errD_real + errD_fake\n        optimizerD.step()\n\n        ############################\n        # (2) Update G network: maximize log(D(G(z)))\n        ###########################\n        netG.zero_grad()\n        labels.fill_(real_label)  # fake labels are real for generator cost\n        output = netD((fake, class_labels))\n        errG = criterion(output, labels)\n        errG.backward()\n        D_G_z2 = output.mean().item()\n        optimizerG.step()\n\n        if step % 500 == 0:\n            print(\n                \"[%d/%d][%d/%d] Loss_D: %.4f Loss_G: %.4f D(x): %.4f D(G(z)): %.4f / %.4f\"\n                % (\n                    epoch + 1,\n                    epochs,\n                    ii,\n                    len(train_loader),\n                    errD.item(),\n                    errG.item(),\n                    D_x,\n                    D_G_z1,\n                    D_G_z2,\n                )\n            )\n\n            valid_image = netG((fixed_noise, class_labels))\n        step += 1\n        lr_schedulerG.step(epoch)\n        lr_schedulerD.step(epoch)\n\n    if epoch % 10 == 0:\n        show_generated_img()\n\n# torch.save(netG.state_dict(), 'generator.pth')\n# torch.save(netD.state_dict(), 'discriminator.pth')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"good_breeds = analyse_generated_by_class(6)\ncreate_submit(good_breeds)","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}