{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import os\nimport random\nimport glob\nimport shutil\nimport warnings\n\nimport cv2\nimport scipy\nimport numpy as np\nimport matplotlib.pyplot as plt\nimport xml.etree.ElementTree as ET\nimport albumentations as A\nfrom sklearn.preprocessing import LabelEncoder\nfrom PIL import Image\nfrom albumentations.pytorch import ToTensor\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.nn.utils import spectral_norm\nfrom torch.optim import Adam\nfrom torch.optim.lr_scheduler import LambdaLR\nfrom torch.utils.data import DataLoader, Dataset\nfrom torchvision.utils import make_grid, save_image\n\n\n%matplotlib inline\nwarnings.filterwarnings('ignore', category=FutureWarning)\ndevice = 'cuda' if torch.cuda.is_available() else 'cpu'","execution_count":null,"outputs":[]},{"metadata":{"_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","collapsed":true,"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","trusted":false},"cell_type":"markdown","source":"# Settings"},{"metadata":{"trusted":true},"cell_type":"code","source":"config = {'DataLoader': {'batch_size': 128,\n                         'shuffle': True},\n          'Generator': {'latent_dim': 120,\n                        'embed_dim': 32,\n                        'ch': 64,\n                        'num_classes': 120,\n                        'use_attn': True},\n          'Discriminator': {'ch': 64,\n                            'num_classes': 120,\n                            'use_attn': True},\n          'sample_latents': {'latent_dim': 120,\n                             'num_classes': 120},\n          'num_iterations': 26000,\n          'decay_start_iteration': 25000,\n          'd_steps': 1,\n          'lr_G': 2e-4,\n          'lr_D': 4e-4,\n          'betas': (0.0, 0.999),\n          'margin': 1.0,\n          'gamma': 0.1,\n          'ema': 0.999,\n          'seed': 42}","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def seed_everything(seed):\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\nseed_everything(config['seed'])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Data Processing"},{"metadata":{},"cell_type":"markdown","source":"## Load Data"},{"metadata":{"trusted":true},"cell_type":"code","source":"root_images = '../input/all-dogs/all-dogs/'\nroot_annots = '../input/annotation/Annotation/'","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"all_files = os.listdir(root_images)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"breeds = glob.glob(root_annots+'*')\nannotations = []\nfor breed in breeds:\n    annotations += glob.glob(breed+'/*')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"breed_map = {}\nfor annotation in annotations:\n    breed = annotation.split('/')[-2]\n    index = breed.split('-')[0]\n    breed_map.setdefault(index, breed)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"all_labels = [breed_map[file.split('_')[0]] for file in all_files]\nle = LabelEncoder()\nall_labels = le.fit_transform(all_labels)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def load_bbox(file):\n    file = str(breed_map[file.split('_')[0]]) + '/' + str(file.split('.')[0])\n    path = os.path.join(root_annots, file)\n    tree = ET.parse(path)\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    \n    return (xmin, ymin, xmax, ymax)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"all_bboxes = [load_bbox(file) for file in all_files]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# make square bounding boxes of original ones,\n# to keep a dog's aspect ratio.\n\ndef get_resized_bbox(height, width, bbox):\n    xmin, ymin, xmax, ymax = bbox\n    xlen = xmax - xmin\n    ylen = ymax - ymin\n    \n    if xlen > ylen:\n        diff = xlen - ylen\n        min_pad = min(ymin, diff//2)\n        max_pad = min(height-ymax, diff-min_pad)\n        ymin = ymin - min_pad\n        ymax = ymax + max_pad\n\n    elif ylen > xlen:\n        diff = ylen - xlen\n        min_pad = min(xmin, diff//2)\n        max_pad = min(width-xmax, diff-min_pad)\n        xmin = xmin - min_pad\n        xmax = xmax + max_pad\n    \n    return xmin, ymin, xmax, ymax","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"resized_bboxes = []\nfor file, bbox in zip(all_files, all_bboxes):\n    img = Image.open(os.path.join(root_images, file))\n    width, height = img.size\n    xmin, ymin, xmax, ymax = get_resized_bbox(height, width, bbox)\n    resized_bboxes.append((xmin, ymin, xmax, ymax))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# crop by square bounding box -> resize and normalize.\n# cv2.INTER_AREA is better than others.\n\ndef load_bboxcrop_resized_image(file, bbox):\n    img = cv2.imread(os.path.join(root_images, file))\n    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n\n    xmin, ymin, xmax, ymax = bbox\n    img = img[ymin:ymax,xmin:xmax]\n\n    transform = A.Compose([A.Resize(64, 64, interpolation=cv2.INTER_AREA),\n                           A.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])\n    img = transform(image=img)['image']\n    \n    return img","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"all_images = [load_bboxcrop_resized_image(f, b) for f, b in zip(all_files, resized_bboxes)]\nall_images = np.array(all_images)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class DogDataset(Dataset):\n    def __init__(self, images, labels):\n        super().__init__()\n        self.images = images\n        self.labels = labels\n        self.transform = A.Compose([A.HorizontalFlip(p=0.5), ToTensor()])\n\n    def __len__(self):\n        return len(self.images)\n\n    def __getitem__(self, idx):\n        img, label = self.images[idx], self.labels[idx]\n        img = self.transform(image=img)['image']\n        label = torch.as_tensor(label, dtype=torch.long)\n\n        return img, label","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# adding uniform noise works well.\n\ndef get_dataiterator(images, labels, dataloader_params, device='cpu'):\n    train_dataset = DogDataset(images, labels)\n    train_dataloader = DataLoader(train_dataset, **dataloader_params)\n    batch_size = dataloader_params['batch_size']\n\n    while True:\n        for imgs, labels in train_dataloader:\n            if batch_size != imgs.size(0):\n                break\n            else:\n                imgs, labels = imgs.to(device), labels.to(device)\n                imgs += (1.0 / 128.0) * torch.rand_like(imgs)\n\n                yield imgs, labels","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# batch size around 64 ~ 128 improves score.\n# ~ 64 are too small, 128 ~ are too large (for 9 hours training). \n\ntrain_dataiterator = get_dataiterator(all_images, all_labels, config['DataLoader'], device=device)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Models"},{"metadata":{"trusted":true},"cell_type":"code","source":"# Attention slightly works.\n\nclass Attention(nn.Module):\n    def __init__(self, channels, reduction_attn=8, reduction_sc=2):\n        super().__init__()\n        self.channles_attn = channels // reduction_attn\n        self.channels_sc = channels // reduction_sc\n        \n        self.conv_query = spectral_norm(nn.Conv2d(channels, self.channles_attn, kernel_size=1, bias=False))\n        self.conv_key = spectral_norm(nn.Conv2d(channels, self.channles_attn, kernel_size=1, bias=False))\n        self.conv_value = spectral_norm(nn.Conv2d(channels, self.channels_sc, kernel_size=1, bias=False))\n        self.conv_attn = spectral_norm(nn.Conv2d(self.channels_sc, channels, kernel_size=1, bias=False))\n        self.gamma = nn.Parameter(torch.zeros(1))\n        \n        nn.init.orthogonal_(self.conv_query.weight.data)\n        nn.init.orthogonal_(self.conv_key.weight.data)\n        nn.init.orthogonal_(self.conv_value.weight.data)\n        nn.init.orthogonal_(self.conv_attn.weight.data)\n\n    def forward(self, x):\n        batch, _, h, w = x.size()\n        \n        proj_query = self.conv_query(x).view(batch, self.channles_attn, -1)\n        proj_key = F.max_pool2d(self.conv_key(x), 2).view(batch, self.channles_attn, -1)\n        \n        attn = torch.bmm(proj_key.permute(0,2,1), proj_query)\n        attn = F.softmax(attn, dim=1)\n        \n        proj_value = F.max_pool2d(self.conv_value(x), 2).view(batch, self.channels_sc, -1)\n        attn = torch.bmm(proj_value, attn)\n        attn = attn.view(batch, self.channels_sc, h, w)\n        attn = self.conv_attn(attn)\n        \n        out = self.gamma * attn + x\n        \n        return out","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Generator"},{"metadata":{"trusted":true},"cell_type":"code","source":"# using label information works well.\n# As for generator, it is realized by conditional batch normalization.\n\nclass CBN2d(nn.Module):\n    def __init__(self, num_features, num_conditions):\n        super().__init__()\n        self.bn = nn.BatchNorm2d(num_features, affine=False)\n        self.embed = spectral_norm(nn.Conv2d(num_conditions, num_features*2, kernel_size=1, bias=False))\n        \n        nn.init.orthogonal_(self.embed.weight.data)\n\n    def forward(self, x, y):\n        out = self.bn(x)\n        embed = self.embed(y.unsqueeze(2).unsqueeze(3))\n        gamma, beta = embed.chunk(2, dim=1)\n        out = (1.0 + gamma) * out + beta \n\n        return out","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# residual block improves convergence speed and generated image's quality.\n# nearest upsampling is better than others.\n\nclass GBlock(nn.Module):\n    def __init__(self, in_channels, out_channels, num_conditions, upsample=False):\n        super().__init__()\n        self.upsample = upsample\n        self.learnable_sc = in_channels != out_channels or upsample\n        \n        self.conv1 = spectral_norm(nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, bias=False))\n        self.conv2 = spectral_norm(nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=False))\n        self.cbn1 = CBN2d(in_channels, num_conditions)\n        self.cbn2 = CBN2d(out_channels, num_conditions)\n        if self.learnable_sc:\n            self.conv_sc = spectral_norm(nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False))\n        self.relu = nn.ReLU()\n\n        nn.init.orthogonal_(self.conv1.weight.data)\n        nn.init.orthogonal_(self.conv2.weight.data)\n        if self.learnable_sc:\n            nn.init.orthogonal_(self.conv_sc.weight.data)\n    \n    def _upsample_conv(self, x, conv):\n        x = F.interpolate(x, scale_factor=2, mode='nearest')\n        x = conv(x)\n        \n        return x\n    \n    def _residual(self, x, y):\n        x = self.relu(self.cbn1(x, y))\n        x = self._upsample_conv(x, self.conv1) if self.upsample else self.conv1(x)\n        x = self.relu(self.cbn2(x, y))\n        x = self.conv2(x)\n        \n        return x\n    \n    def _shortcut(self, x):\n        if self.learnable_sc:\n            x = self._upsample_conv(x, self.conv_sc) if self.upsample else self.conv_sc(x)\n            \n        return x\n    \n    def forward(self, x, y):\n        return self._shortcut(x) + self._residual(x, y)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# shared embedding of class labels, and hierarchical latent noise, work well.\n# this architecture is the same as BigGAN except for channel size.\n\nclass Generator(nn.Module):\n    def __init__(self, latent_dim, ch, num_classes, embed_dim, use_attn=False):\n        super().__init__()\n        self.latent_dim = latent_dim\n        self.ch = ch\n        self.num_classes = num_classes\n        self.embed_dim = embed_dim\n        self.use_attn = use_attn\n        self.num_chunk = 5\n        num_latents = self.__get_num_latents()\n        \n        self.embed = nn.Embedding(num_classes, embed_dim)\n        self.fc = spectral_norm(nn.Linear(num_latents[0], ch*8*4*4, bias=False))\n        self.block1 = GBlock(ch*8, ch*8, num_latents[1], upsample=True)\n        self.block2 = GBlock(ch*8, ch*4, num_latents[2], upsample=True)\n        self.block3 = GBlock(ch*4, ch*2, num_latents[3], upsample=True)\n        if use_attn:\n            self.attn = Attention(ch*2)\n        self.block4 = GBlock(ch*2, ch, num_latents[4], upsample=True)\n        self.bn = nn.BatchNorm2d(ch)\n        self.relu = nn.ReLU()\n        self.conv_last = spectral_norm(nn.Conv2d(ch, 3, kernel_size=3, padding=1, bias=False))\n        self.tanh = nn.Tanh()\n        \n        nn.init.orthogonal_(self.embed.weight.data)\n        nn.init.orthogonal_(self.fc.weight.data)\n        nn.init.orthogonal_(self.conv_last.weight.data)\n        nn.init.constant_(self.bn.weight.data, 1.0)\n        nn.init.constant_(self.bn.bias.data, 0.0)\n    \n    def __get_num_latents(self):\n        xs = torch.empty(self.latent_dim).chunk(self.num_chunk)\n        num_latents = [x.size(0) for x in xs]\n        for i in range(1, self.num_chunk):\n            num_latents[i] += self.embed_dim\n        \n        return num_latents\n    \n    def forward(self, x, y):\n        xs = x.chunk(self.num_chunk, dim=1)\n        y = self.embed(y)\n        \n        h = self.fc(xs[0])\n        h = h.view(h.size(0), self.ch*8, 4, 4)\n        h = self.block1(h, torch.cat([y, xs[1]], dim=1))\n        h = self.block2(h, torch.cat([y, xs[2]], dim=1))\n        h = self.block3(h, torch.cat([y, xs[3]], dim=1))\n        if self.use_attn:\n            h = self.attn(h)\n        h = self.block4(h, torch.cat([y, xs[4]], dim=1))\n        h = self.relu(self.bn(h))\n        out = self.tanh(self.conv_last(h))\n        \n        return out","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Discriminator"},{"metadata":{"trusted":true},"cell_type":"code","source":"# residual block improves convergence speed and generated image's quality.\n\nclass DBlock(nn.Module):\n    def __init__(self, in_channels, out_channels, downsample=False, optimized=False):\n        super().__init__()\n        self.downsample = downsample\n        self.optimized = optimized\n        self.learnable_sc = in_channels != out_channels or downsample\n        \n        self.conv1 = spectral_norm(nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, bias=False))\n        self.conv2 = spectral_norm(nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=False))\n        if self.learnable_sc:\n            self.conv_sc = spectral_norm(nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False))\n        self.relu = nn.ReLU()\n        \n        nn.init.orthogonal_(self.conv1.weight.data)\n        nn.init.orthogonal_(self.conv2.weight.data)\n        if self.learnable_sc:\n            nn.init.orthogonal_(self.conv_sc.weight.data)\n\n    def _residual(self, x):\n        if not self.optimized:\n            x = self.relu(x)\n        x = self.conv1(x)\n        x = self.relu(x)\n        x = self.conv2(x)\n        if self.downsample:\n            x = F.avg_pool2d(x, 2)\n        \n        return x\n    \n    def _shortcut(self, x):\n        if self.learnable_sc:\n            if self.optimized:\n                x = self.conv_sc(F.avg_pool2d(x, 2)) if self.downsample else self.conv_sc(x)\n            else:\n                x = F.avg_pool2d(self.conv_sc(x), 2) if self.downsample else self.conv_sc(x)\n        \n        return x\n    \n    def forward(self, x):\n        return self._shortcut(x) + self._residual(x)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# this architecture is the altered version of BigGAN Discriminator.\n# <- using residual block, projection.\n\n# but those points are different from original.\n# - reduce channel size.\n# - reduce model depth (remove last residual block).\n# - add minibatch stddev.\n# - with auxiliary classifier (ACGAN).\n#   <- improve image's quality and stabilize training.\n\nclass Discriminator(nn.Module):\n    def __init__(self, ch, num_classes, use_attn=False):\n        super().__init__()\n        self.ch = ch\n        self.num_classes = num_classes\n        self.use_attn = use_attn\n        \n        self.block1 = DBlock(3, ch, downsample=True, optimized=True)\n        if use_attn:\n            self.attn = Attention(ch)\n        self.block2 = DBlock(ch, ch*2, downsample=True)\n        self.block3 = DBlock(ch*2, ch*4, downsample=True)\n        self.block4 = DBlock(ch*4, ch*8, downsample=True)\n        self.block5 = DBlock(ch*8+1, ch*8, downsample=False)\n        self.relu = nn.ReLU()\n        self.fc = spectral_norm(nn.Linear(ch*8, 1, bias=False))\n        self.embed = spectral_norm(nn.Embedding(num_classes, ch*8))\n        self.clf = spectral_norm(nn.Linear(ch*8, num_classes, bias=False))\n        \n        nn.init.orthogonal_(self.fc.weight.data)\n        nn.init.orthogonal_(self.embed.weight.data)\n        nn.init.orthogonal_(self.clf.weight.data)\n    \n    def minibatch_stddev(self, x, group_size=4, eps=1e-8):\n        shape = x.size()\n        y = x.view(group_size, -1, shape[1], shape[2], shape[3])\n        y -= torch.mean(y, dim=0, keepdim=True)\n        y = torch.mean(y.pow(2), dim=0)\n        y = torch.sqrt(y + eps)\n        y = torch.mean(y, dim=[1,2,3], keepdim=True)\n        y = y.repeat(group_size, 1, shape[2], shape[3])\n\n        return torch.cat([x, y], dim=1)\n    \n    def forward(self, x, y):\n        h = self.block1(x)\n        if self.use_attn:\n            h = self.attn(h)\n        h = self.block2(h)\n        h = self.block3(h)\n        h = self.block4(h)\n        h = self.block5(self.minibatch_stddev(h))\n        h = self.relu(h)\n        h = torch.sum(h, dim=(2,3))\n        \n        out = self.fc(h)\n        out += torch.sum(self.embed(y)*h, dim=1, keepdim=True)\n        \n        ac = self.clf(h)\n        ac = F.log_softmax(ac, dim=1)\n        \n        return out, ac","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Train GANs"},{"metadata":{"trusted":true},"cell_type":"code","source":"netG = Generator(**config['Generator']).to(device, torch.float32)\nnetD = Discriminator(**config['Discriminator']).to(device, torch.float32)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Exponential moving average of generator weights works well.\n\nnetGE = Generator(**config['Generator']).to(device, torch.float32)\nnetGE.load_state_dict(netG.state_dict());","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"optim_G = Adam(params=netG.parameters(), lr=config['lr_G'], betas=config['betas'])\noptim_D = Adam(params=netD.parameters(), lr=config['lr_D'], betas=config['betas'])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"decay_iter = config['num_iterations'] - config['decay_start_iteration']\nif decay_iter > 0:\n    lr_lambda_G = lambda x: (max(0,1-x/decay_iter))\n    lr_lambda_D = lambda x: (max(0,1-x/(decay_iter*config['d_steps'])))\n    lr_sche_G = LambdaLR(optim_G, lr_lambda=lr_lambda_G)\n    lr_sche_D = LambdaLR(optim_D, lr_lambda=lr_lambda_D)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def calc_advloss_D(real, fake, margin=1.0):\n    loss_real = torch.mean((real - fake.mean() - margin) ** 2)\n    loss_fake = torch.mean((fake - real.mean() + margin) ** 2)\n    loss = (loss_real + loss_fake) / 2\n    \n    return loss","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def calc_advloss_G(real, fake, margin=1.0):\n    loss_real = torch.mean((real - fake.mean() + margin) ** 2)\n    loss_fake = torch.mean((fake - real.mean() - margin) ** 2)\n    loss = (loss_real + loss_fake) / 2\n    \n    return loss","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# auxiliary classifier loss.\n# this loss weighted by gamma (0.1) is added to adversarial loss.\n# coefficient gamma is quite sensitive.\n\ncriterion = nn.NLLLoss().to(device, torch.float32)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def sample_latents(batch_size, latent_dim, num_classes):\n    latents = torch.randn((batch_size, latent_dim), dtype=torch.float32, device=device)\n    labels = torch.randint(0, num_classes, size=(batch_size,), dtype=torch.long, device=device)\n    \n    return latents, labels","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"step = 1\n\nwhile True:\n    # Discriminator\n    for i in range(config['d_steps']):\n        for param in netD.parameters():\n            param.requires_grad_(True)\n    \n        optim_D.zero_grad()\n\n        real_imgs, real_labels = train_dataiterator.__next__()\n        batch_size = real_imgs.size(0)\n\n        latents, fake_labels = sample_latents(batch_size, **config['sample_latents'])\n        fake_imgs = netG(latents, fake_labels).detach()\n        \n        preds_real, preds_real_labels = netD(real_imgs, real_labels)\n        preds_fake, _ = netD(fake_imgs, fake_labels)\n\n        loss_D = calc_advloss_D(preds_real, preds_fake, config['margin'])\n        loss_D += config['gamma'] * criterion(preds_real_labels, real_labels)\n        loss_D.backward()\n        optim_D.step()\n        \n        if (decay_iter > 0) and (step > config['decay_start_iteration']):\n            lr_sche_D.step()\n\n    # Generator\n    for param in netD.parameters():\n        param.requires_grad_(False)\n\n    optim_G.zero_grad()\n    \n    real_imgs, real_labels = train_dataiterator.__next__()\n    batch_size = real_imgs.size(0)\n    \n    latents, fake_labels = sample_latents(batch_size, **config['sample_latents'])\n    fake_imgs = netG(latents, fake_labels)\n\n    preds_real, _ = netD(real_imgs, real_labels)\n    preds_fake, preds_fake_labels = netD(fake_imgs, fake_labels)\n\n    loss_G = calc_advloss_G(preds_real, preds_fake, config['margin'])\n    loss_G += config['gamma'] * criterion(preds_fake_labels, fake_labels)\n    loss_G.backward()\n    optim_G.step()\n    \n    if (decay_iter > 0) and (step > config['decay_start_iteration']):\n        lr_sche_G.step()\n    \n    # Update Generator Eval\n    for param_G, param_GE in zip(netG.parameters(), netGE.parameters()):\n        param_GE.data.mul_(config['ema']).add_((1-config['ema'])*param_G.data)\n    for buffer_G, buffer_GE in zip(netG.buffers(), netGE.buffers()):\n        buffer_GE.data.mul_(config['ema']).add_((1-config['ema'])*buffer_G.data)\n            \n    # stopping\n    if step < config['num_iterations']:\n        step += 1\n    else:\n        print('total step: {}'.format(step))\n        break","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Generate Samples"},{"metadata":{"trusted":true},"cell_type":"code","source":"def truncated_normal(size, threshold=2.0, dtype=torch.float32, device='cpu'):\n    x = scipy.stats.truncnorm.rvs(-threshold, threshold, size=size)\n    x = torch.from_numpy(x).to(device, dtype)\n\n    return x","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def generate_eval_samples(generator, batch_size, latent_dim, num_classes):\n    latents = truncated_normal((batch_size, latent_dim), dtype=torch.float32, device=device)\n    labels =  torch.randint(0, num_classes, size=(batch_size,), dtype=torch.long, device=device)\n    \n    with torch.no_grad():\n        imgs = (generator(latents, labels) + 1) / 2\n    \n    return imgs","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def make_submissions(generator, user_images_unzipped_path, latent_dim, num_classes):\n    if not os.path.exists(user_images_unzipped_path):\n        os.mkdir(user_images_unzipped_path)\n    \n    sample_batch_size = 50\n    num_samples = 10000\n    \n    for i in range(0, num_samples, sample_batch_size):\n        imgs = generate_eval_samples(generator, sample_batch_size, latent_dim, num_classes)\n        for j, img in enumerate(imgs):\n            save_image(img, os.path.join(user_images_unzipped_path, f'image_{i+j:05d}.png'))\n    \n    shutil.make_archive('images', 'zip', user_images_unzipped_path)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"user_images_unzipped_path = '../output_images'\nmake_submissions(netGE, user_images_unzipped_path, **config['sample_latents'])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# See some examples"},{"metadata":{"trusted":true},"cell_type":"code","source":"repeat_breeds = 5\nnrow = 15\nsplit_size = 50\n\nlatent_dim = config['sample_latents']['latent_dim']\nnum_classes = config['sample_latents']['num_classes']\n\nall_labels = torch.arange(num_classes, dtype=torch.long, device=device)\nall_labels = all_labels.repeat_interleave(repeat_breeds)\nlabels_split = all_labels.split(split_size)\n\nimgs_list = []\nfor labels in labels_split:\n    batch_size = labels.size(0)\n    latents = truncated_normal((batch_size, latent_dim), threshold=1.5, dtype=torch.float32, device=device)\n    with torch.no_grad():\n        imgs = (netGE(latents, labels) + 1) / 2\n        imgs_list.append(imgs)\n    \nall_imgs = torch.cat(imgs_list, dim=0)\nall_imgs = make_grid(all_imgs, nrow=nrow, normalize=False)\nall_imgs = all_imgs.mul_(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to('cpu', torch.uint8).numpy()\n\nplt.figure(figsize=(2*nrow, 2*(num_classes*repeat_breeds)//nrow))\nplt.imshow(all_imgs);","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.6.4","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"}},"nbformat":4,"nbformat_minor":1}