{"cells":[{"metadata":{"_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","trusted":true},"cell_type":"code","source":"import matplotlib.pyplot as plt\nimport numpy as np\nimport torch\nfrom torch import nn, optim\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader \nfrom torchvision import datasets, transforms\nfrom torchvision.utils import save_image, make_grid\nimport os\nimport torchvision\nimport xml.etree.ElementTree as ET\nfrom tqdm import tqdm_notebook as tqdm\nimport time\nimport torch.nn.utils.spectral_norm as SN\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n# PATH = '../input/generative-dog-images/all-dogs/'\n# ANNOTATION_PATH = \"../input/generative-dog-images/annotation/Annotation/\"\nPATH = '../input/all-dogs/'\nANNOTATION_PATH = \"../input/annotation/Annotation/\"\nOUTPUT_PATH = '../output_images'\n\nIMG_DIR = \"output_images/\"\n\nIMAGE_SIZE = 64\nBATCH_SIZE = 32\nPRINT_EVERY=1000\nSHOW_EVERY=1000\nPLOT_EVERY=200\nEPOCHS=200\n\nlrD = 0.001\nlrG = 0.001\nbeta1 = 0.5\nbeta2 = 0.999\n\nnz = 512\nSPECTRAL_NORM=True\nNORMALIZATION='selfmod' # selfmod or adain\nRANDOM_NOISE=True\nPIXEL_NORM=True\n\nLOSS='NS' #NS or WGAN\nUSE_SOFT_NOISY_LABELS = True\nINVERT_LABELS = True\nreal_label = 0.9\nfake_label = 0\n\nN_IMAGES=10000\nComputeLB = False","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Data loading"},{"metadata":{"trusted":true},"cell_type":"code","source":"class DogsDataset(Dataset):\n    def __init__(self, root, transform1=None, transform2=None):\n        self.root = root\n#         super().__init__(root, transform1=transform1, transform2=transform2)\n        self.transform1 = transform1\n        self.transform2 = transform2\n        \n        self.samples, self.crops = self._load_subfolders_images(self.root)\n        if len(self.samples) == 0:\n            raise RuntimeError(\"Found 0 files in subfolders of: {}\".format(self.root))\n            \n    def _load_subfolders_images(self, root):\n        IMG_EXTENSIONS = ('.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif', '.tiff', '.webp')\n\n        def is_valid_file(x):\n            return torchvision.datasets.folder.has_file_allowed_extension(x, IMG_EXTENSIONS)\n\n        imgs = []\n        img_paths = []\n        bboxes = []\n\n        paths = []\n        for root, _, fnames in sorted(os.walk(root)):\n            for fname in sorted(fnames):\n                path = os.path.join(root, fname)\n                paths.append(path)\n\n        pbar = tqdm(paths, desc='Loading cropped images')\n\n        for path in pbar:\n            if is_valid_file(path):\n                # Load image\n                img = torchvision.datasets.folder.default_loader(path)\n\n                # Get bounding boxes\n                annotation_basename = os.path.splitext(os.path.basename(path))[0]\n                annotation_dirname = next(dirname for dirname in os.listdir(ANNOTATION_PATH) if dirname.startswith(annotation_basename.split('_')[0]))\n                annotation_filename = os.path.join(ANNOTATION_PATH, annotation_dirname, annotation_basename)\n                tree = ET.parse(annotation_filename)\n                root = tree.getroot()\n                size = root.find('size')\n                width = int(size.find('width').text)\n                height = int(size.find('height').text)\n                objects = root.findall('object')\n                objects = root.findall('object')\n                for o in objects:    \n                    bndbox = o.find('bndbox') \n                \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\n                    xmin = max(0, xmin - 4)        # 4 : margin\n                    xmax = min(width, xmax + 4)\n                    ymin = max(0, ymin - 4)\n                    ymax = min(height, ymax + 4)\n\n                    w = np.min((xmax - xmin, ymax - ymin))\n                    w = min(w, width, height)                     # available w\n\n                    if w > xmax - xmin:\n                        xmin = min(max(0, xmin - int((w - (xmax - xmin))/2)), width - w)\n                        xmax = xmin + w\n                    if w > ymax - ymin:\n                        ymin = min(max(0, ymin - int((w - (ymax - ymin))/2)), height - w)\n                        ymax = ymin + w\n\n                    img_ = img.crop((xmin, ymin, xmax, ymax)) #img[ymin:ymin+w, xmin:xmin+w, :]      # [h,w,c]\n\n                    if np.mean(img_) != 0:\n                        img = img_\n                        \n#                     if self.transform1 is not None:\n#                         img = self.transform1(img)\n                        \n#                     imgs.append(img)\n                        img_paths.append(path)\n                        bboxes.append((xmin, ymin, xmax, ymax))\n                \n                pbar.set_postfix_str(\"{} cropped images loaded\".format(len(img_paths)))\n\n        return img_paths, bboxes\n    \n    def __getitem__(self, index):\n        path = self.samples[index]\n        sample = torchvision.datasets.folder.default_loader(path)\n        sample = sample.crop(self.crops[index])\n        \n        if self.transform2 is not None:\n            sample = self.transform2(sample)\n\n        return sample\n\n    def __len__(self):\n        return len(self.samples)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# First preprocessing of data\ntransform1 = None #transforms.Compose([transforms.RandomResizedCrop(64, (0.85, 1.0))])\n\n\n# Data augmentation and converting to tensors\ntransform2 = transforms.Compose([transforms.RandomResizedCrop(64, (0.85, 1.0)),\n                                 transforms.ColorJitter(),\n                                 transforms.RandomHorizontalFlip(p=0.5), \n                                 transforms.ToTensor(),\n                                 transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])\n\n\ntrain_data = DogsDataset(\n    PATH,\n    transform1,\n    transform2\n)\n\ntrain_loader = torch.utils.data.DataLoader(train_data, shuffle=True,\n                                           batch_size=BATCH_SIZE, num_workers=4)\n\n# Plot some training images\n# real_batch = next(iter(train_loader))\n# plt.figure(figsize=(8,8))\n# plt.axis(\"off\")\n# plt.title(\"Training Images\")\n# plt.imshow(np.transpose(make_grid(real_batch.to(device)[:64], padding=2, normalize=True).cpu(),(1,2,0)))","execution_count":null,"outputs":[]},{"metadata":{"_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","trusted":true},"cell_type":"code","source":"def slerp(val, low, high):\n    low_norm = low/torch.norm(low, dim=1, keepdim=True)\n    high_norm = high/torch.norm(high, dim=1, keepdim=True)\n    omega = torch.acos((low_norm*high_norm).sum(1))\n    so = torch.sin(omega)\n    res = (torch.sin((1.0-val)*omega)/so).unsqueeze(1)*low + (torch.sin(val*omega)/so).unsqueeze(1) * high\n    return res\n\nclass BatchNormModulate2d(nn.Module):\n    \"\"\"\n    Similar to batch norm, but with learnable weights and bias\n    \"\"\"\n    def __init__(self, num_features, dim_in, eps=2e-5, momentum=0.1, affine=True,\n                 track_running_stats=True, use_sn=True):\n        super().__init__()\n        self.num_features = num_features\n        self.dim_in = dim_in\n        self.bn = nn.BatchNorm2d(num_features, affine=False)\n        self.gamma = nn.Sequential(\n            nn.Linear(dim_in, num_features, bias=True),\n            nn.ReLU(),\n            nn.Linear(num_features, num_features, bias=False)\n        )\n        self.beta = nn.Sequential(\n            nn.Linear(dim_in, num_features, bias=True),\n            nn.ReLU(),\n            nn.Linear(num_features, num_features, bias=False)\n        )\n\n    def forward(self, x, z):\n        out = self.bn(x)\n        gamma = self.gamma(z)\n        beta = self.beta(z)\n        out = gamma.view(-1, self.num_features, 1, 1) * out + beta.view(-1, self.num_features, 1, 1)\n        return out\n    \nclass PixelNorm(nn.Module):\n    def __init__(self, epsilon=1e-8):\n        \"\"\"\n            @notice: avoid in-place ops.\n            https://discuss.pytorch.org/t/encounter-the-runtimeerror-one-of-the-variables-needed-for-gradient-computation-has-been-modified-by-an-inplace-operation/836/3\n        \"\"\"\n        super(PixelNorm, self).__init__()\n        self.epsilon = epsilon\n\n    def forward(self, x):\n        tmp  = torch.mul(x, x) # or x ** 2\n        tmp1 = torch.rsqrt(torch.mean(tmp, dim=1, keepdim=True) + self.epsilon)\n\n        return x * tmp1\n    \nclass GaussianNoise(nn.Module):\n    def __init__(self, channels):\n        super().__init__()\n        self.weight = nn.Parameter(torch.zeros(1, channels, 1, 1))\n        \n    def forward(self, x, noise=None):\n        if noise is None:\n            noise = torch.randn(x.shape[0], 1, x.shape[2], x.shape[3], device=x.device, dtype=x.dtype)\n        return x + self.weight.view(1, -1, 1, 1) * noise.to(x.device)\n    \nclass AdaIn(nn.Module):\n    \"\"\"\n    latent_dim represents dimension of latent vector similar to style vector in StyleGAN\n    \"\"\"\n    def __init__(self, in_channel, latent_dim):\n        super().__init__()\n\n        self.norm = nn.InstanceNorm2d(in_channel)\n        self.style = nn.Linear(latent_dim, in_channel * 2)\n\n    def forward(self, input, style):\n        style = self.style(style).unsqueeze(2).unsqueeze(3)\n        gamma, beta = style.chunk(2, 1)\n\n        out = self.norm(input)\n        out = gamma * out + beta\n\n        return out\n        \nclass Projection(nn.Module):\n    def __init__(self, \n                 nz, \n                 in_channel, \n                 out_channel, \n                 shape, \n                 bias=False, \n                 spectral_norm=False, \n                 normalization='selfmod', \n                 random_noise=False,\n                use_pixel_norm=False):\n        super().__init__()\n        self.shape = shape\n        self.linear = nn.Linear(in_channel, out_channel, bias=bias)\n        self.conv = nn.Conv2d(shape[0], shape[0], 3, 1, 1, bias=bias)\n        if spectral_norm:\n            self.linear = SN(self.linear)\n            self.conv = SN(self.conv)\n            \n        self.noise1 = None\n        if random_noise:\n            self.noise1 = GaussianNoise(shape[0])\n            self.noise2 = GaussianNoise(shape[0])\n            \n        self.pixel_norm = None\n        if use_pixel_norm:\n            self.pixel_norm = PixelNorm() \n            \n        if normalization == 'adain':\n            self.norm1 = AdaIn(shape[0], nz)\n            self.norm2 = AdaIn(shape[0], nz)\n        else:\n            self.norm1 = BatchNormModulate2d(shape[0], nz)\n            self.norm2 = BatchNormModulate2d(shape[0], nz)\n            \n            \n    def forward(self, x, nz):\n        x = self.linear(x)\n        x = x.view([x.shape[0]] + self.shape)\n        if self.noise1 is not None:\n            x = self.noise1(x)\n        x = F.leaky_relu(x, 0.2)\n        if self.pixel_norm is not None:\n            x = self.pixel_norm(x)\n        x = self.norm1(x, nz)\n        x = self.conv(x)\n        if self.noise2 is not None:\n            x = self.noise2(x)\n        x = F.leaky_relu(x, 0.2)\n        if self.pixel_norm is not None:\n            x = self.pixel_norm(x)\n        x = self.norm2(x, nz)\n        return x\n    \nclass UpConvBlock(nn.Module):\n    \"\"\"\n    normalization is 'selfmod', 'adain'\n    \"\"\"\n    def __init__(self, \n                 nz, \n                 in_channel, \n                 out_channel, \n                 kernel=4, \n                 stride=2,\n                 padding=1, \n                 bias=False, \n                 spectral_norm=False, \n                 normalization='selfmod', \n                 random_noise=False,\n                use_pixel_norm=False):\n        super().__init__()\n        self.conv1 = nn.ConvTranspose2d(in_channel, out_channel, kernel, stride, padding, bias=bias)\n        self.conv2 = nn.Conv2d(out_channel, out_channel, 3, 1, 1, bias=bias)\n        if spectral_norm:\n            self.conv1 = SN(self.conv1)\n            self.conv2 = SN(self.conv2)\n            \n        self.noise1 = None\n        self.noise2 = None\n        if random_noise:\n            self.noise1 = GaussianNoise(out_channel)\n            self.noise2 = GaussianNoise(out_channel)\n            \n        self.pixel_norm = None\n        if use_pixel_norm:\n            self.pixel_norm = PixelNorm() \n        \n        if normalization == 'adain':\n            self.norm1 = AdaIn(out_channel, nz)\n            self.norm2 = AdaIn(out_channel, nz)\n        else:\n            self.norm1 = BatchNormModulate2d(out_channel, nz)\n            self.norm2 = BatchNormModulate2d(out_channel, nz)\n            \n        \n    def forward(self, x, noise):\n        x = self.conv1(x)\n        if self.noise1 is not None:\n            x = self.noise1(x)\n        x = F.leaky_relu(x, 0.2)\n        if self.pixel_norm is not None:\n            x = self.pixel_norm(x)\n        x = self.norm1(x, noise)\n        x = self.conv2(x)\n        if self.noise2 is not None:\n            x = self.noise2(x)\n        x = F.leaky_relu(x, 0.2)\n        if self.pixel_norm is not None:\n            x = self.pixel_norm(x)\n        x = self.norm2(x, noise)\n        return x\n        \n\nclass Generator(nn.Module):\n    def __init__(self, nz, nfeats, nchannels, spectral_norm=False, normalization='selfmod', random_noise=False, use_pixel_norm=False):\n        super(Generator, self).__init__()\n        d = nfeats*8\n        \n        self.linear = Projection(nz, nz, 8*8*d, [512, 8, 8], False, spectral_norm, normalization, random_noise, use_pixel_norm) \n        # state size. (nfeats*8) x 4 x 4\n        \n        self.conv2 = UpConvBlock(nz, nfeats*8, nfeats*4, 4, 2, 1, False, spectral_norm, normalization, random_noise, use_pixel_norm) \n        # state size. (nfeats*8) x 8 x 8\n        \n        self.conv3 = UpConvBlock(nz, nfeats*4, nfeats*2, 4, 2, 1, False, spectral_norm, normalization, random_noise, use_pixel_norm) \n        # state size. (nfeats*4) x 16 x 16\n        \n        self.conv4 = UpConvBlock(nz, nfeats*2, nfeats, 4, 2, 1, False, spectral_norm, normalization, random_noise, use_pixel_norm)\n        # state size. (nfeats * 2) x 32 x 32\n        \n        self.conv5 = nn.ConvTranspose2d(nfeats, nchannels, 3, 1, 1, bias=False)\n        if spectral_norm:\n            self.conv5 = SN(self.conv5)\n        # state size. (nfeats) x 64 x 64\n\n    def forward(self, x):\n        out = self.linear(x.view(x.shape[0],-1), x)\n        out = self.conv2(out, x)\n        out = self.conv3(out, x)\n        out = self.conv4(out, x)\n        out = torch.tanh(self.conv5(out))\n        \n        return out\n\n\nclass Discriminator(nn.Module):\n    def __init__(self, nchannels, nfeats):\n        super(Discriminator, self).__init__()\n        d = nfeats\n        # input is (nchannels) x 64 x 64\n        \n        self.model = nn.Sequential(\n            SN(nn.Conv2d(nchannels, d, 3, 1, 1, bias=False)),\n            nn.LeakyReLU(0.2, inplace=True),\n\n            SN(nn.Conv2d(d, d * 2, 4, 2, 1, bias=False)),\n            nn.LeakyReLU(0.2, inplace=True),\n            SN(nn.Conv2d(d*2, d*2, 3, 1, 1, bias=False)),\n            nn.LeakyReLU(0.2, inplace=True),\n\n            SN(nn.Conv2d(d*2, d*4, 4, 2, 1, bias=False)),\n            nn.LeakyReLU(0.2, inplace=True),\n            SN(nn.Conv2d(d*4, d*4, 3, 1, 1, bias=False)),\n            nn.LeakyReLU(0.2, inplace=True),\n\n            SN(nn.Conv2d(d*4, d*8, 4, 2, 1, bias=False)),\n            nn.LeakyReLU(0.2, inplace=True),\n            SN(nn.Conv2d(d*8, d*8, 3, 1, 1, bias=False)),\n            nn.LeakyReLU(0.2, inplace=True)\n        )\n        \n        self.linear = SN(nn.Linear(8*8*d*8, 1, bias=False))\n        \n        \n    def forward(self, x):\n        out = self.model(x)\n        out = out.view(out.shape[0], -1)\n        out = self.linear(out)\n        return out","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Define training loop"},{"metadata":{"trusted":true},"cell_type":"code","source":"def gradient_penalty(x, y, f):\n    # interpolation\n    shape = [x.size(0)] + [1] * (x.dim() - 1)\n    alpha = torch.rand(shape).to(x.device)\n    z = x + alpha * (y - x)\n\n    # gradient penalty\n    z = Variable(z, requires_grad=True).to(x.device)\n    o = f(z)\n    g = grad(o, z, grad_outputs=torch.ones(o.size()).to(z.device), create_graph=True)[0].view(z.size(0), -1)\n    gp = ((g.norm(p=2, dim=1) - 1)**2).mean()\n    return gp\n    \ndef R1Penalty(real_img, f):\n    # gradient penalty\n    reals = Variable(real_img, requires_grad=True).to(real_img.device)\n    real_logit = f(reals)\n    apply_loss_scaling = lambda x: x * torch.exp(x * torch.Tensor([np.float32(np.log(2.0))]).to(real_img.device))\n    undo_loss_scaling = lambda x: x * torch.exp(-x * torch.Tensor([np.float32(np.log(2.0))]).to(real_img.device))\n\n    real_logit = apply_loss_scaling(torch.sum(real_logit))\n    real_grads = grad(real_logit, reals, grad_outputs=torch.ones(real_logit.size()).to(reals.device), create_graph=True)[0].view(reals.size(0), -1)\n    real_grads = undo_loss_scaling(real_grads)\n    r1_penalty = torch.sum(torch.mul(real_grads, real_grads))\n    return r1_penalty\n\ndef G_wgan(G, D, nz, batch_size):\n    noise = torch.randn(batch_size, nz, device=device)\n    fake_images = G(noise)\n    fake_logit = D(fake_images)\n    G_loss = -fake_logit.mean()\n    return G_loss\n\ndef D_wgan_gp(G, D, real_images, nz, lammy=10.0, eps=0.001):\n    batch_size = real_images.shape[0]\n    real_logit = D(real_images)\n    noise = torch.randn(batch_size, nz, device=device)\n    fake_images = G(noise)\n    fake_logit = D(fake_images.detach())\n    D_loss = fake_logit.mean() - real_logit.mean()\n    D_loss = D_loss + gradient_penalty(real_images.data, fake_images.data, D) * lammy\n#     D_loss += real_logit.mean()**2 * eps\n    return D_loss, real_logit.mean(), fake_logit.mean()\n    \ndef D_NS(G, D, real_images, nz, lammy=0.0):\n    batch_size = real_images.shape[0]\n    real_logit = D(real_images)\n    D_loss_real = F.softplus(-real_logit).mean() \n    noise = torch.randn(batch_size, nz, device=device)\n    fake_images = G(noise)\n    fake_logit = D(fake_images.detach())\n    D_loss_fake = F.softplus(fake_logit).mean() \n    D_loss = D_loss_real + D_loss_fake\n    if lammy != 0.0:\n        r1_penalty = R1Penalty(real_images.detach(), D)\n        D_loss = D_loss + r1_penalty * (lammy * 0.5)\n    return D_loss, D_loss_real, D_loss_fake\n    \ndef G_NS(G, D, nz, batch_size):\n    noise = torch.randn(batch_size, nz, device=device)\n    fake_images = G(noise)\n    fake_logit = D(fake_images)\n    G_loss = F.softplus(-fake_logit).mean() \n    return G_loss\n\ndef D_BCE_NS(G, D, real_images, nz, real_labels, fake_labels, lammy=0.0):\n    batch_size = real_images.shape[0]\n    real_logit = D(real_images)\n    D_loss_real = F.binary_cross_entropy_with_logits(real_logit, real_labels)\n    noise = torch.randn(batch_size, nz, device=device)\n    fake_images = G(noise)\n    fake_logit = D(fake_images.detach())\n    D_loss_fake = F.binary_cross_entropy_with_logits(fake_logit, fake_labels)\n    D_loss = D_loss_real + D_loss_fake\n    if lammy != 0.0:\n        r1_penalty = R1Penalty(real_images.detach(), D)\n        D_loss = D_loss + r1_penalty * (lammy * 0.5)\n    return D_loss, D_loss_real, D_loss_fake\n    \ndef G_BCE_NS(G, D, nz, batch_size, real_labels):\n    noise = torch.randn(batch_size, nz, device=device)\n    fake_images = G(noise)\n    fake_logit = D(fake_images)\n    G_loss = F.binary_cross_entropy_with_logits(fake_logit, real_labels)\n    return G_loss\n\nclass Trainer:\n    def __init__(self, nz, G, D, lammy=0.0, track_grads=False):\n        self.nz = nz\n        self.track_grads=track_grads\n        self.G = G\n        self.D = D\n        self.fixed_noise = torch.randn(64, self.nz)\n        self.lammy = lammy\n        \n        self.d_losses = []\n        self.d_losses_real = []\n        self.d_losses_fake = []\n        self.g_losses = []\n        self.img_list = []\n        self.g_grads = []\n        self.d_grads = []\n        \n    def check_grads(self, model):\n        grads = []\n        for n, p in model.named_parameters():\n            if not p.grad is None and p.requires_grad and \"bias\" not in n:\n                grads.append(float(p.grad.abs().mean()))\n        return grads\n            \n        \n    def train(self, epochs, loader, criterion, optim_G, optim_D, scheduler_D, scheduler_G, print_every, show_every, plot_every, loss='NS'):\n        step = 0\n        \n        for epoch in tqdm(range(epochs)):\n            for ii, real_images in enumerate(loader):\n                batch_size = real_images.size(0)\n                \n                # Train Discriminator\n                self.D.zero_grad()\n                real_images = real_images.to(device)\n                \n                if loss == 'WGAN':\n                    D_loss, Dx, D_Gz = D_wgan_gp(self.G, self.D, real_images, self.nz)\n                elif loss == 'BCE':\n                    if USE_SOFT_NOISY_LABELS:\n                        real_labels = torch.empty((batch_size, 1), device=device).uniform_(0.80, 0.95)\n                        fake_labels = torch.empty((batch_size, 1), device=device).uniform_(0.05, 0.20)\n                    else:\n                        real_labels = torch.full((batch_size, 1), 0.95, device=device)\n                        fake_labels = torch.full((batch_size, 1), 0.05, device=device)\n\n                    if INVERT_LABELS and random.random() < 0.01:\n                        real_labels, fake_labels = fake_labels, real_labels\n                    D_loss, Dx, D_Gz = D_BCE_NS(self.G, self.D, real_images, self.nz, real_labels, fake_labels, self.lammy)\n                else:\n                    D_loss, Dx, D_Gz = D_NS(self.G, self.D, real_images, self.nz, self.lammy)\n                \n                D_loss.backward()\n                optim_D.step()\n                \n                # Train Generator\n                self.G.zero_grad()\n                \n                if loss == 'WGAN':\n                    G_loss = G_wgan(self.G, self.D, self.nz, batch_size)\n                elif loss == 'BCE':\n                    G_loss = G_BCE_NS(self.G, self.D, self.nz, batch_size, real_labels)\n                else:\n                    G_loss = G_NS(self.G, self.D, self.nz, batch_size)\n                G_loss.backward()\n                optim_G.step()\n                \n                if step % print_every == 0:\n                    print('[%d/%d][%d/%d]\\tLoss_D: %.4f\\tLoss_G: %.4f\\tD(x): %.4f\\tD(G(z)): %.4f' \n                      % (epoch, epochs, ii, len(loader),\n                         D_loss.item(), G_loss.item(), Dx, D_Gz))\n                \n                if step % plot_every == 0:\n                    self.d_losses.append(D_loss.item())\n                    self.d_losses_real.append(Dx)\n                    self.d_losses_fake.append(D_Gz)\n                    self.g_losses.append(G_loss.item())\n                    if self.track_grads:\n                        self.g_grads.append(self.check_grads(self.G))\n                        self.d_grads.append(self.check_grads(self.D))\n                \n#                 if (step % show_every == 0) or ((epoch == epochs-1) and (ii == len(loader)-1)):\n#                     with torch.no_grad():\n#                         fake = self.G(self.fixed_noise.to(device)).detach().cpu()\n#                     self.img_list.append(make_grid(fake, padding=2, normalize=True))\n                \n                step += 1\n            scheduler_D.step()\n            scheduler_G.step()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"netG = Generator(nz, 64, 3, SPECTRAL_NORM, NORMALIZATION, RANDOM_NOISE, PIXEL_NORM).to(device)\n# print(netG)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"netD = Discriminator(3, 64).to(device)\n# print(netD)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def weights_init(m):\n    if type(m) == nn.Linear or type(m) == nn.Conv2d or type(m) == nn.ConvTranspose2d:\n        torch.nn.init.kaiming_uniform(m.weight)\n        if m.bias is not None:\n            m.bias.data.zero_()\n#             m.bias.data.fill_(0.01)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\nnetG.apply(weights_init)\nnetD.apply(weights_init)\n\ncriterion = nn.BCELoss()\n\noptimizerD = optim.Adam(netD.parameters(), lr=lrD, betas=(beta1, beta2))\noptimizerG = optim.Adam(netG.parameters(), lr=lrG, betas=(beta1, beta2))\nscheduler_D = optim.lr_scheduler.ExponentialLR(optimizerD, gamma=0.99)\nscheduler_G = optim.lr_scheduler.ExponentialLR(optimizerG, gamma=0.99)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"trainer = Trainer(nz, netG, netD, track_grads=True)\ntrainer.train(EPOCHS, train_loader, criterion, optimizerG, optimizerD, scheduler_D, scheduler_G, PRINT_EVERY, SHOW_EVERY, PLOT_EVERY, loss=LOSS)\n               \n# torch.save(netG.state_dict(), 'generator.pth')\n# torch.save(netD.state_dict(), 'discriminator.pth')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Look at results"},{"metadata":{"trusted":true},"cell_type":"code","source":"plt.figure(figsize=(10,5))\nplt.title(\"Generator and Discriminator Loss During Training\")\nplt.plot(trainer.g_losses,label=\"G\")\nplt.plot(trainer.d_losses,label=\"D\")\nplt.xlabel(\"iterations\")\nplt.ylabel(\"Loss\")\nplt.legend()\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Visualize progression"},{"metadata":{"trusted":true},"cell_type":"code","source":"# import matplotlib.animation as animation\n# from IPython.display import HTML\n\n# fig = plt.figure(figsize=(8,8))\n# plt.axis(\"off\")\n# ims = [[plt.imshow(np.transpose(i,(1,2,0)), animated=False)] for i in trainer.img_list]\n# ani = animation.ArtistAnimation(fig, ims, interval=1000, repeat_delay=1000, blit=True)\n# HTML(ani.to_jshtml())\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Real versus fake images"},{"metadata":{"trusted":true},"cell_type":"code","source":"# # Grab a batch of real images from the dataloader\n\n# # # Plot the real images\n# plt.figure(figsize=(15,15))\n\n# # Plot the fake images from the last epoch\n# plt.subplot(1,2,1)\n# plt.axis(\"off\")\n# plt.title(\"Fake Images\")\n# plt.imshow(np.transpose(trainer.img_list[-1],(1,2,0)))\n# plt.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Check the gradients for G\nfor i in trainer.g_grads: plt.plot(i)\nplt.legend(range(5))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Check the gradients for D\nfor i in trainer.d_grads: plt.plot(i)\nplt.legend(range(5))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"if not os.path.exists(OUTPUT_PATH):\n    os.mkdir(OUTPUT_PATH)\nwith torch.no_grad():\n    for i_batch in range(0, N_IMAGES, 25):\n        gen_z = torch.randn(25, nz, device=device)\n        gen_images = netG(gen_z)\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, :, :, :]*0.5 + 0.5, os.path.join(OUTPUT_PATH, f'image_{i_batch+i_image:05d}.png'))\n\nimport shutil\nshutil.make_archive('images', 'zip', OUTPUT_PATH)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Compute leadboard"},{"metadata":{"trusted":true},"cell_type":"code","source":"from __future__ import absolute_import, division, print_function\nimport numpy as np\nimport os\nimport gzip, pickle\nimport tensorflow as tf\nfrom scipy import linalg\nimport pathlib\nimport urllib\nimport warnings\nfrom tqdm import tqdm\nfrom PIL import Image\n\nclass KernelEvalException(Exception):\n    pass\n\nmodel_params = {\n    'Inception': {\n        'name': 'Inception', \n        'imsize': 64,\n        'output_layer': 'Pretrained_Net/pool_3:0', \n        'input_layer': 'Pretrained_Net/ExpandDims:0',\n        'output_shape': 2048,\n        'cosine_distance_eps': 0.1\n        }\n}\n\ndef create_model_graph(pth):\n    \"\"\"Creates a graph from saved GraphDef file.\"\"\"\n    # Creates graph from saved graph_def.pb.\n    with tf.gfile.FastGFile( pth, 'rb') as f:\n        graph_def = tf.GraphDef()\n        graph_def.ParseFromString( f.read())\n        _ = tf.import_graph_def( graph_def, name='Pretrained_Net')\n\ndef _get_model_layer(sess, model_name):\n    # layername = 'Pretrained_Net/final_layer/Mean:0'\n    layername = model_params[model_name]['output_layer']\n    layer = sess.graph.get_tensor_by_name(layername)\n    ops = layer.graph.get_operations()\n    for op_idx, op in enumerate(ops):\n        for o in op.outputs:\n            shape = o.get_shape()\n            if shape._dims != []:\n              shape = [s.value for s in shape]\n              new_shape = []\n              for j, s in enumerate(shape):\n                if s == 1 and j == 0:\n                  new_shape.append(None)\n                else:\n                  new_shape.append(s)\n              o.__dict__['_shape_val'] = tf.TensorShape(new_shape)\n    return layer\n\ndef get_activations(images, sess, model_name, batch_size=50, verbose=False):\n    \"\"\"Calculates the activations of the pool_3 layer for all images.\n\n    Params:\n    -- images      : Numpy array of dimension (n_images, hi, wi, 3). The values\n                     must lie between 0 and 256.\n    -- sess        : current session\n    -- batch_size  : the images numpy array is split into batches with batch size\n                     batch_size. A reasonable batch size depends on the disposable hardware.\n    -- verbose    : If set to True and parameter out_step is given, the number of calculated\n                     batches is reported.\n    Returns:\n    -- A numpy array of dimension (num images, 2048) that contains the\n       activations of the given tensor when feeding inception with the query tensor.\n    \"\"\"\n    inception_layer = _get_model_layer(sess, model_name)\n    n_images = images.shape[0]\n    if batch_size > n_images:\n        print(\"warning: batch size is bigger than the data size. setting batch size to data size\")\n        batch_size = n_images\n    n_batches = n_images//batch_size + 1\n    pred_arr = np.empty((n_images,model_params[model_name]['output_shape']))\n    for i in tqdm(range(n_batches)):\n        if verbose:\n            print(\"\\rPropagating batch %d/%d\" % (i+1, n_batches), end=\"\", flush=True)\n        start = i*batch_size\n        if start+batch_size < n_images:\n            end = start+batch_size\n        else:\n            end = n_images\n                    \n        batch = images[start:end]\n        pred = sess.run(inception_layer, {model_params[model_name]['input_layer']: batch})\n        pred_arr[start:end] = pred.reshape(-1,model_params[model_name]['output_shape'])\n    if verbose:\n        print(\" done\")\n    return pred_arr\n\n\n# def calculate_memorization_distance(features1, features2):\n#     neigh = NearestNeighbors(n_neighbors=1, algorithm='kd_tree', metric='euclidean')\n#     neigh.fit(features2) \n#     d, _ = neigh.kneighbors(features1, return_distance=True)\n#     print('d.shape=',d.shape)\n#     return np.mean(d)\n\ndef normalize_rows(x: np.ndarray):\n    \"\"\"\n    function that normalizes each row of the matrix x to have unit length.\n\n    Args:\n     ``x``: A numpy matrix of shape (n, m)\n\n    Returns:\n     ``x``: The normalized (by row) numpy matrix.\n    \"\"\"\n    return np.nan_to_num(x/np.linalg.norm(x, ord=2, axis=1, keepdims=True))\n\n\ndef cosine_distance(features1, features2):\n    # print('rows of zeros in features1 = ',sum(np.sum(features1, axis=1) == 0))\n    # print('rows of zeros in features2 = ',sum(np.sum(features2, axis=1) == 0))\n    features1_nozero = features1[np.sum(features1, axis=1) != 0]\n    features2_nozero = features2[np.sum(features2, axis=1) != 0]\n    norm_f1 = normalize_rows(features1_nozero)\n    norm_f2 = normalize_rows(features2_nozero)\n\n    d = 1.0-np.abs(np.matmul(norm_f1, norm_f2.T))\n    print('d.shape=',d.shape)\n    print('np.min(d, axis=1).shape=',np.min(d, axis=1).shape)\n    mean_min_d = np.mean(np.min(d, axis=1))\n    print('distance=',mean_min_d)\n    return mean_min_d\n\n\ndef distance_thresholding(d, eps):\n    if d < eps:\n        return d\n    else:\n        return 1\n\ndef calculate_frechet_distance(mu1, sigma1, mu2, sigma2, eps=1e-6):\n    \"\"\"Numpy implementation of the Frechet Distance.\n    The Frechet distance between two multivariate Gaussians X_1 ~ N(mu_1, C_1)\n    and X_2 ~ N(mu_2, C_2) is\n            d^2 = ||mu_1 - mu_2||^2 + Tr(C_1 + C_2 - 2*sqrt(C_1*C_2)).\n            \n    Stable version by Dougal J. Sutherland.\n\n    Params:\n    -- mu1 : Numpy array containing the activations of the pool_3 layer of the\n             inception net ( like returned by the function 'get_predictions')\n             for generated samples.\n    -- mu2   : The sample mean over activations of the pool_3 layer, precalcualted\n               on an representive data set.\n    -- sigma1: The covariance matrix over activations of the pool_3 layer for\n               generated samples.\n    -- sigma2: The covariance matrix over activations of the pool_3 layer,\n               precalcualted on an representive data set.\n\n    Returns:\n    --   : The Frechet Distance.\n    \"\"\"\n\n    mu1 = np.atleast_1d(mu1)\n    mu2 = np.atleast_1d(mu2)\n\n    sigma1 = np.atleast_2d(sigma1)\n    sigma2 = np.atleast_2d(sigma2)\n\n    assert mu1.shape == mu2.shape, \"Training and test mean vectors have different lengths\"\n    assert sigma1.shape == sigma2.shape, \"Training and test covariances have different dimensions\"\n\n    diff = mu1 - mu2\n\n    # product might be almost singular\n    covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False)\n    if not np.isfinite(covmean).all():\n        msg = \"fid calculation produces singular product; adding %s to diagonal of cov estimates\" % eps\n        warnings.warn(msg)\n        offset = np.eye(sigma1.shape[0]) * eps\n        # covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset))\n        covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset))\n    \n    # numerical error might give slight imaginary component\n    if np.iscomplexobj(covmean):\n        if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3):\n            m = np.max(np.abs(covmean.imag))\n            raise ValueError(\"Imaginary component {}\".format(m))\n        covmean = covmean.real\n\n    # covmean = tf.linalg.sqrtm(tf.linalg.matmul(sigma1,sigma2))\n\n    print('covmean.shape=',covmean.shape)\n    # tr_covmean = tf.linalg.trace(covmean)\n\n    tr_covmean = np.trace(covmean)\n    return diff.dot(diff) + np.trace(sigma1) + np.trace(sigma2) - 2 * tr_covmean\n    # return diff.dot(diff) + tf.linalg.trace(sigma1) + tf.linalg.trace(sigma2) - 2 * tr_covmean\n#-------------------------------------------------------------------------------\n\n\ndef calculate_activation_statistics(images, sess, model_name, batch_size=50, verbose=False):\n    \"\"\"Calculation of the statistics used by the FID.\n    Params:\n    -- images      : Numpy array of dimension (n_images, hi, wi, 3). The values\n                     must lie between 0 and 255.\n    -- sess        : current session\n    -- batch_size  : the images numpy array is split into batches with batch size\n                     batch_size. A reasonable batch size depends on the available hardware.\n    -- verbose     : If set to True and parameter out_step is given, the number of calculated\n                     batches is reported.\n    Returns:\n    -- mu    : The mean over samples of the activations of the pool_3 layer of\n               the incption model.\n    -- sigma : The covariance matrix of the activations of the pool_3 layer of\n               the incption model.\n    \"\"\"\n    act = get_activations(images, sess, model_name, batch_size, verbose)\n    mu = np.mean(act, axis=0)\n    sigma = np.cov(act, rowvar=False)\n    return mu, sigma, act\n    \ndef _handle_path_memorization(path, sess, model_name, is_checksize, is_check_png):\n    path = pathlib.Path(path)\n    files = list(path.glob('*.jpg')) + list(path.glob('*.png'))\n    imsize = model_params[model_name]['imsize']\n\n    # In production we don't resize input images. This is just for demo purpose. \n    x = np.array([np.array(img_read_checks(fn, imsize, is_checksize, imsize, is_check_png)) for fn in files])\n    m, s, features = calculate_activation_statistics(x, sess, model_name)\n    del x #clean up memory\n    return m, s, features\n\n# check for image size\ndef img_read_checks(filename, resize_to, is_checksize=False, check_imsize = 64, is_check_png = False):\n    im = Image.open(str(filename))\n    if is_checksize and im.size != (check_imsize,check_imsize):\n        raise KernelEvalException('The images are not of size '+str(check_imsize))\n    \n    if is_check_png and im.format != 'PNG':\n        raise KernelEvalException('Only PNG images should be submitted.')\n\n    if resize_to is None:\n        return im\n    else:\n        return im.resize((resize_to,resize_to),Image.ANTIALIAS)\n\ndef calculate_kid_given_paths(paths, model_name, model_path, feature_path=None, mm=[], ss=[], ff=[]):\n    ''' Calculates the KID of two paths. '''\n    tf.reset_default_graph()\n    create_model_graph(str(model_path))\n    with tf.Session() as sess:\n        sess.run(tf.global_variables_initializer())\n        m1, s1, features1 = _handle_path_memorization(paths[0], sess, model_name, is_checksize = True, is_check_png = True)\n        if len(mm) != 0:\n            m2 = mm\n            s2 = ss\n            features2 = ff\n        elif feature_path is None:\n            m2, s2, features2 = _handle_path_memorization(paths[1], sess, model_name, is_checksize = False, is_check_png = False)\n        else:\n            with np.load(feature_path) as f:\n                m2, s2, features2 = f['m'], f['s'], f['features']\n\n        print('m1,m2 shape=',(m1.shape,m2.shape),'s1,s2=',(s1.shape,s2.shape))\n        print('starting calculating FID')\n        fid_value = calculate_frechet_distance(m1, s1, m2, s2)\n        print('done with FID, starting distance calculation')\n        distance = cosine_distance(features1, features2)        \n        return fid_value, distance, m2, s2, features2","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import zipfile\nif ComputeLB:\n  \n    # UNCOMPRESS OUR IMGAES\n    with zipfile.ZipFile(\"../working/images.zip\",\"r\") as z:\n        z.extractall(\"../tmp/images2/\")\n\n    # COMPUTE LB SCORE\n    m2 = []; s2 =[]; f2 = []\n    user_images_unzipped_path = '../tmp/images2/'\n    images_path = [user_images_unzipped_path,'../input/generative-dog-images/all-dogs/all-dogs/']\n    public_path = '../input/dog-face-generation-competition-kid-metric-input/classify_image_graph_def.pb'\n\n    fid_epsilon = 10e-15\n\n    fid_value_public, distance_public, m2, s2, f2 = calculate_kid_given_paths(images_path, 'Inception', public_path, mm=m2, ss=s2, ff=f2)\n    distance_public = distance_thresholding(distance_public, model_params['Inception']['cosine_distance_eps'])\n    print(\"FID_public: \", fid_value_public, \"distance_public: \", distance_public, \"multiplied_public: \",\n            fid_value_public /(distance_public + fid_epsilon))\n    \n    # REMOVE FILES TO PREVENT KERNEL ERROR OF TOO MANY FILES\n    ! rm -r ../tmp","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}