{"cells":[{"metadata":{},"cell_type":"markdown","source":"I have to modify the architecture (decrease number of channels, number of samples per phase) to deal with the limited time.\nThe source for styleGAN:\n[https://github.com/rosinality/style-based-gan-pytorch](https://github.com/rosinality/style-based-gan-pytorch)\n\nThe dataloader:\n[https://www.kaggle.com/speedwagon/ram-dataloader](https://www.kaggle.com/speedwagon/ram-dataloader)"},{"metadata":{"trusted":true},"cell_type":"code","source":"import argparse\nimport random\nimport math\n\nfrom tqdm import tqdm\nimport numpy as np\nfrom PIL import Image\n\nimport torch\nfrom torch import nn, optim\nfrom torch.nn import functional as F\nfrom torch.autograd import Variable, grad\nfrom torch.utils.data import DataLoader\nfrom torchvision import datasets, transforms, utils\nimport torchvision.utils as vutils\nfrom torch.utils.data import Dataset, DataLoader\nimport os\nfrom os.path import join\nimport xml.etree.ElementTree as ET\nimport matplotlib.pyplot as plt\n\ndef requires_grad(model, flag=True):\n    for p in model.parameters():\n        p.requires_grad = flag\n\n\ndef accumulate(model1, model2, decay=0.999):\n    par1 = dict(model1.named_parameters())\n    par2 = dict(model2.named_parameters())\n\n    for k in par1.keys():\n        par1[k].data.mul_(decay).add_(1 - decay, par2[k].data)\n\n\ndef adjust_lr(optimizer, lr):\n    for group in optimizer.param_groups:\n        mult = group.get('mult', 1)\n        group['lr'] = lr * mult\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class DogDataset(Dataset):\n    def __init__(self, img_dir, annotations_dir, transform1=None, transform2=None):\n\n        self.img_dir = img_dir\n        self.img_names = os.listdir(img_dir)\n        self.transform1 = transform1\n        self.transform2 = transform2\n        \n        self.imgs = []\n        for img_name in self.img_names:\n            path = join(img_dir, img_name)\n            img = datasets.folder.default_loader(path)\n    \n            # Crop image\n            annotation_basename = os.path.splitext(os.path.basename(path))[0]\n            annotation_dirname = next(dirname for dirname in os.listdir(annotations_dir) if dirname.startswith(annotation_basename.split('_')[0]))\n            annotation_filename = os.path.join(annotations_dir, annotation_dirname, annotation_basename)\n            tree = ET.parse(annotation_filename)\n            root = tree.getroot()\n            objects = root.findall('object')\n            \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                bbox = (xmin, ymin, xmax, ymax)\n                img_ = img.crop(bbox)\n                # Some crop's are black. if crop is black then don't crop\n                if np.mean(img_) != 0:\n                    img = img_\n\n                if self.transform1 is not None:\n                    img = self.transform1(img)\n\n                self.imgs.append(img)\n\n    def __getitem__(self, index):\n        img = self.imgs[index]\n        \n        if self.transform2 is not None:\n            img = self.transform2(img)\n        \n        return img\n\n    def __len__(self):\n        return len(self.imgs)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def sample_data(PATH_IMG, batch_size, image_size=32):\n    transform1 = transforms.Compose([transforms.Resize(image_size),\n                                    transforms.CenterCrop(image_size)])\n\n    # Data augmentation and converting to tensors\n    random_transforms = [transforms.RandomRotation(degrees=5)]\n    transform2 = transforms.Compose([transforms.RandomHorizontalFlip(p=0.5),\n                                     transforms.RandomApply(random_transforms, p=0.3), \n                                     transforms.ToTensor(),\n                                     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])\n\n    train_dataset = DogDataset(img_dir='../input/all-dogs/all-dogs/',\n                               annotations_dir='../input/annotation/Annotation/',\n                               transform1=transform1,\n                               transform2=transform2)\n\n    train_loader = DataLoader(dataset=train_dataset,\n                              batch_size=batch_size,\n                              shuffle=True,\n                              num_workers=4)\n\n    return train_loader, train_dataset","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch\n\nfrom torch import nn\nfrom torch.nn import init\nfrom torch.nn import functional as F\nfrom torch.autograd import Variable\n\nfrom math import sqrt\n\nimport random\n\nn_class_age = 6\nn_repeat = 1\n\ndef init_linear(linear):\n    init.xavier_normal(linear.weight)\n    linear.bias.data.zero_()\n\n\ndef init_conv(conv, glu=True):\n    init.kaiming_normal(conv.weight)\n    if conv.bias is not None:\n        conv.bias.data.zero_()\n\n\nclass EqualLR:\n    def __init__(self, name):\n        self.name = name\n\n    def compute_weight(self, module):\n        weight = getattr(module, self.name + '_orig')\n        fan_in = weight.data.size(1) * weight.data[0][0].numel()\n\n        return weight * sqrt(2 / fan_in)\n\n    @staticmethod\n    def apply(module, name):\n        fn = EqualLR(name)\n\n        weight = getattr(module, name)\n        del module._parameters[name]\n        module.register_parameter(name + '_orig', nn.Parameter(weight.data))\n        module.register_forward_pre_hook(fn)\n\n        return fn\n\n    def __call__(self, module, input):\n        weight = self.compute_weight(module)\n        setattr(module, self.name, weight)\n\n\ndef equal_lr(module, name='weight'):\n    EqualLR.apply(module, name)\n\n    return module\n\n\nclass PixelNorm(nn.Module):\n    def __init__(self):\n        super().__init__()\n\n    def forward(self, input):\n        return input / torch.sqrt(torch.mean(input ** 2, dim=1, keepdim=True) + 1e-8)\n\n\nclass Blur(nn.Module):\n    def __init__(self):\n        super().__init__()\n\n        weight = torch.tensor([[1, 2, 1], [2, 4, 2], [1, 2, 1]], dtype=torch.float32)\n        weight = weight.view(1, 1, 3, 3)\n        weight = weight / weight.sum()\n        self.register_buffer('weight', weight)\n\n    def forward(self, input):\n        return F.conv2d(\n            input,\n            self.weight.repeat(input.shape[1], 1, 1, 1),\n            padding=1,\n            groups=input.shape[1],\n        )\n\n\nclass EqualConv2d(nn.Module):\n    def __init__(self, *args, **kwargs):\n        super().__init__()\n\n        conv = nn.Conv2d(*args, **kwargs)\n        conv.weight.data.normal_()\n        conv.bias.data.zero_()\n        self.conv = equal_lr(conv)\n\n    def forward(self, input):\n        return self.conv(input)\n\n\nclass EqualLinear(nn.Module):\n    def __init__(self, in_dim, out_dim):\n        super().__init__()\n\n        linear = nn.Linear(in_dim, out_dim)\n        linear.weight.data.normal_()\n        linear.bias.data.zero_()\n\n        self.linear = equal_lr(linear)\n\n    def forward(self, input):\n        return self.linear(input)\n\n\nclass ConvBlock(nn.Module):\n    def __init__(\n        self,\n        in_channel,\n        out_channel,\n        kernel_size,\n        padding,\n        kernel_size2=None,\n        padding2=None,\n        pixel_norm=True,\n        spectral_norm=False,\n    ):\n        super().__init__()\n\n        pad1 = padding\n        pad2 = padding\n        if padding2 is not None:\n            pad2 = padding2\n\n        kernel1 = kernel_size\n        kernel2 = kernel_size\n        if kernel_size2 is not None:\n            kernel2 = kernel_size2\n\n        self.conv = nn.Sequential(\n            EqualConv2d(in_channel, out_channel, kernel1, padding=pad1),\n            nn.LeakyReLU(0.2),\n            EqualConv2d(out_channel, out_channel, kernel2, padding=pad2),\n            nn.LeakyReLU(0.2),\n        )\n\n    def forward(self, input):\n        out = self.conv(input)\n\n        return out\n\n\nclass AdaptiveInstanceNorm(nn.Module):\n    def __init__(self, in_channel, style_dim):\n        super().__init__()\n\n        self.norm = nn.InstanceNorm2d(in_channel)\n        self.style = EqualLinear(style_dim, in_channel * 2)\n\n        self.style.linear.bias.data[:in_channel] = 1\n        self.style.linear.bias.data[in_channel:] = 0\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\n\nclass NoiseInjection(nn.Module):\n    def __init__(self, channel):\n        super().__init__()\n\n        self.weight = nn.Parameter(torch.zeros(1, channel, 1, 1))\n\n    def forward(self, image, noise):\n        return image + self.weight * noise\n\n\nclass ConstantInput(nn.Module):\n    def __init__(self, channel, size=8):\n        super().__init__()\n\n        self.input = nn.Parameter(torch.randn(1, channel, size, size))\n\n    def forward(self, input):\n        batch = input.shape[0]\n        out = self.input.repeat(batch, 1, 1, 1)\n\n        return out\n\n\nclass StyledConvBlock(nn.Module):\n    def __init__(\n        self,\n        in_channel,\n        out_channel,\n        kernel_size=3,\n        padding=1,\n        style_dim=512,\n        initial=False,\n    ):\n        super().__init__()\n\n        if initial:\n            self.conv1 = ConstantInput(in_channel)\n\n        else:\n            self.conv1 = EqualConv2d(\n                in_channel, out_channel, kernel_size, padding=padding\n            )\n\n        self.noise1 = equal_lr(NoiseInjection(out_channel))\n        self.adain1 = AdaptiveInstanceNorm(out_channel, style_dim)\n        self.lrelu1 = nn.LeakyReLU(0.2)\n\n        self.conv2 = EqualConv2d(out_channel, out_channel, kernel_size, padding=padding)\n        self.noise2 = equal_lr(NoiseInjection(out_channel))\n        self.adain2 = AdaptiveInstanceNorm(out_channel, style_dim)\n        self.lrelu2 = nn.LeakyReLU(0.2)\n\n    def forward(self, input, style, noise):\n        out = self.conv1(input)\n        out = self.noise1(out, noise)\n        out = self.adain1(out, style)\n        out = self.lrelu1(out)\n\n        out = self.conv2(out)\n        out = self.noise2(out, noise)\n        out = self.adain2(out, style)\n        out = self.lrelu2(out)\n\n        return out\n\n\nclass Generator(nn.Module):\n    def __init__(self, code_dim):\n        super().__init__()\n\n        self.progression = nn.ModuleList(\n            [\n                StyledConvBlock(512, 512, 3, 1, initial=True),  #8\n#                 StyledConvBlock(512, 512, 3, 1),\n                StyledConvBlock(512, 256, 3, 1),                #16\n                StyledConvBlock(256, 128, 3, 1),                #32\n                StyledConvBlock(128, 64, 3, 1),                 #64\n                StyledConvBlock(64, 32, 3, 1),\n                StyledConvBlock(128, 64, 3, 1),\n                StyledConvBlock(64, 32, 3, 1),\n                StyledConvBlock(32, 16, 3, 1),\n            ]\n        )\n\n        self.to_rgb = nn.ModuleList(\n            [\n                EqualConv2d(512, 3, 1),\n#                 EqualConv2d(512, 3, 1),\n                EqualConv2d(256, 3, 1),\n                EqualConv2d(128, 3, 1),\n                EqualConv2d(64, 3, 1),\n                EqualConv2d(32, 3, 1),\n                EqualConv2d(64, 3, 1),\n                EqualConv2d(32, 3, 1),\n                EqualConv2d(16, 3, 1),\n            ]\n        )\n\n        # self.blur = Blur()\n\n    def forward(self, style, noise, step=0, alpha=-1, mixing_range=(-1, -1)):\n        out = noise[0]\n\n        if len(style) < 2:\n            inject_index = [len(self.progression) + 1]\n\n        else:\n            inject_index = random.sample(list(range(step)), len(style) - 1)\n\n        crossover = 0\n\n        for i, (conv, to_rgb) in enumerate(zip(self.progression, self.to_rgb)):\n            if mixing_range == (-1, -1):\n                if crossover < len(inject_index) and i > inject_index[crossover]:\n                    crossover = min(crossover + 1, len(style))\n\n                style_step = style[crossover]\n\n            else:\n                if mixing_range[0] <= i <= mixing_range[1]:\n                    style_step = style[1]\n\n                else:\n                    style_step = style[0]\n\n            if i > 0 and step > 0:\n                upsample = F.interpolate(\n                    out, scale_factor=2, mode='bilinear', align_corners=False\n                )\n                out = conv(upsample, style_step, noise[i])\n\n            else:\n                out = conv(out, style_step, noise[i])\n\n            if i == step:\n                out = to_rgb(out)\n\n                if i > 0 and 0 <= alpha < 1:\n                    skip_rgb = self.to_rgb[i - 1](upsample)\n                    out = (1 - alpha) * skip_rgb + alpha * out\n\n                break\n\n        return out\n\n\nclass StyledGenerator(nn.Module):\n    def __init__(self, code_dim=512, n_mlp=8):\n        super().__init__()\n\n        self.generator = Generator(code_dim)\n\n        layers = [PixelNorm()]\n        for i in range(n_mlp):\n            layers.append(EqualLinear(code_dim, code_dim))\n            layers.append(nn.LeakyReLU(0.2))\n\n        self.style = nn.Sequential(*layers)\n\n    def forward(\n        self,\n        input,\n        noise=None,\n        step=0,\n        alpha=-1,\n        mean_style=None,\n        style_weight=0,\n        mixing_range=(-1, -1),\n    ):\n        styles = []\n        if type(input) not in (list, tuple):\n            input = [input]\n\n        for i in input:\n            styles.append(self.style(i))\n\n        batch = input[0].shape[0]\n\n        if noise is None:\n            noise = []\n\n            for i in range(step + 1):\n                size = 8 * 2 ** i\n                noise.append(torch.randn(batch, 1, size, size, device=input[0].device))\n\n        if mean_style is not None:\n            styles_norm = []\n\n            for style in styles:\n                styles_norm.append(mean_style + style_weight * (style - mean_style))\n\n            styles = styles_norm\n\n        return self.generator(styles, noise, step, alpha, mixing_range=mixing_range)\n\n    def mean_style(self, input):\n        style = self.style(input).mean(0, keepdim=True)\n\n        return style\n\n\nclass Discriminator(nn.Module):\n    def __init__(self):\n        super().__init__()\n\n        self.progression = nn.ModuleList(\n            [\n                ConvBlock(16, 32, 3, 1),\n                ConvBlock(32, 64, 3, 1),\n                ConvBlock(64, 128, 3, 1),\n                ConvBlock(32, 64, 3, 1),\n                ConvBlock(64, 128, 3, 1),\n                ConvBlock(128, 128, 3, 1),\n                ConvBlock(128, 256, 3, 1),\n#                 ConvBlock(512, 512, 3, 1),\n                ConvBlock(257, 512, 3, 1, 8, 0),\n            ]\n        )\n\n        self.from_rgb_ = nn.ModuleList(\n            [\n                EqualConv2d(3, 16, 1),\n                EqualConv2d(3, 32, 1),\n                EqualConv2d(3, 64, 1),\n                EqualConv2d(3, 32, 1),\n                EqualConv2d(3, 64, 1),\n                EqualConv2d(3, 128, 1),\n                EqualConv2d(3, 128, 1),\n#                 EqualConv2d(3+n_class_age, 512, 1),\n                EqualConv2d(3, 256, 1),\n            ]\n        )\n\n        self.n_layer = len(self.progression)\n\n        self.linear = EqualLinear(512, 1)\n\n    def forward(self, input, step=0, alpha=-1):\n        \n        img_size = input.shape[2]\n\n        for i in range(step, -1, -1):\n            index = self.n_layer - i - 1\n\n            if i == step:\n                input_ = input\n                out = self.from_rgb_[index](input_)\n\n            if i == 0:\n                out_std = torch.sqrt(out.var(0, unbiased=False) + 1e-8)\n                mean_std = out_std.mean()\n                mean_std = mean_std.expand(out.size(0), 1, 8, 8)\n                out = torch.cat([out, mean_std], 1)\n\n            out = self.progression[index](out)\n\n            if i > 0:\n                # out = F.avg_pool2d(out, 2)\n                out = F.interpolate(\n                    out, scale_factor=0.5, mode='bilinear', align_corners=False\n                )\n\n                if i == step and 0 <= alpha < 1:\n                    skip_rgb = F.interpolate(\n                        input_, scale_factor=0.5, mode='bilinear', align_corners=False\n                    )\n                    skip_rgb = self.from_rgb_[index + 1](skip_rgb)\n\n                    out = (1 - alpha) * skip_rgb + alpha * out\n\n        out = out.squeeze(2).squeeze(2)\n#         print(input.size(), out.size(), step)\n        out = self.linear(out)\n\n        return out","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"code_size = 512\nbatch_size = 16\nn_critic = 1\nIs_Show = False\nclass Args:\n    n_gpu = 4\n#     phase = 600_000\n    phase = 150_000\n    lr = 0.001\n    init_size = 8\n    max_size = 64\n    mixing = False\n    loss = 'wgan-gp'\n    data = 'folder'\n    path = '/home/quang/working/Dog_kaggle/data/dog-resize/all-dogs/'\n    sched = None\n    \nargs = Args()\ngenerator = nn.DataParallel(StyledGenerator(code_size)).cuda()\ndiscriminator = nn.DataParallel(Discriminator()).cuda()\n\nclass_loss = nn.CrossEntropyLoss()\ng_optimizer = optim.Adam(\n    generator.module.generator.parameters(), lr=args.lr, betas=(0.0, 0.99)\n)    \ng_optimizer.add_param_group(\n    {\n        'params': generator.module.style.parameters(),\n        'lr': args.lr * 0.01,\n        'mult': 0.01,\n    }\n)\n\nd_optimizer = optim.Adam(discriminator.parameters(), lr=args.lr, betas=(0.0, 0.99))\n\nif args.sched:\n    args.lr = {128: 0.0015, 256: 0.002, 512: 0.003, 1024: 0.003}\n    args.batch = {4: 512, 8: 256, 16: 128, 32: 64, 64: 32, 128: 32, 256: 32}\n\nelse:\n    args.lr = {128: 0.0015, 256: 0.002, 512: 0.003, 1024: 0.003}\n    args.batch = {4: 32, 8: 32, 16: 32, 32: 32, 64: 32, 128: 16, 256: 8}\n\nargs.gen_sample = {512: (8, 4), 1024: (4, 2)}\n\nargs.batch_default = 32","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"step = int(math.log2(args.init_size)) - 3\nresolution = 8 * 2 ** step\nloader, dog_dataset = sample_data(\n    args.path, args.batch.get(resolution, args.batch_default), resolution\n)\ndata_loader = iter(loader)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"adjust_lr(g_optimizer, args.lr.get(resolution, 0.001))\nadjust_lr(d_optimizer, args.lr.get(resolution, 0.001))\n\npbar = tqdm(range(60_000))\n# pbar = tqdm(range(100))\n\nrequires_grad(generator, False)\nrequires_grad(discriminator, True)\n\ndisc_loss_val = 0\ngen_loss_val = 0\ngrad_loss_val = 0\n\nalpha = 0\nused_sample = 0\n\nn_repeat = 2\n\nprint ('Resolution: ', resolution, '|Step: ', step, '|Batch_size: ', args.batch.get(resolution, args.batch_default), ' |Generator lr: ', \n      g_optimizer.state_dict()['param_groups'][0]['lr'], ' |Style lr: ', g_optimizer.state_dict()['param_groups'][1]['lr'])\nfor i in pbar:\n    discriminator.zero_grad()\n\n    alpha = min(1, 1 / args.phase * (used_sample + 1))\n\n    if used_sample > args.phase * 2 and step < (int(math.log2(args.max_size)) - 3):\n        step += 1\n\n        if step > int(math.log2(args.max_size)) - 3:\n            step = int(math.log2(args.max_size)) - 3\n\n        else:\n            alpha = 0\n            used_sample = 0\n\n        resolution = 8 * 2 ** step\n        del loader\n        del dog_dataset\n        loader, dog_dataset = sample_data(\n            args.path, args.batch.get(resolution, args.batch_default), resolution\n        )\n        data_loader = iter(loader)\n\n        adjust_lr(g_optimizer, args.lr.get(resolution, 0.001))\n        adjust_lr(d_optimizer, args.lr.get(resolution, 0.001))\n        print ('Resolution: ', resolution, '|Step: ', step, '|Batch_size: ', args.batch.get(resolution, args.batch_default), ' |Generator lr: ', \n              g_optimizer.state_dict()['param_groups'][0]['lr'], ' |Style lr: ', g_optimizer.state_dict()['param_groups'][1]['lr'])\n\n    try:\n        real_image = next(data_loader)\n\n    except (OSError, StopIteration):\n        data_loader = iter(loader)\n        real_image = next(data_loader)\n\n    used_sample += real_image.shape[0]\n\n    b_size = real_image.size(0)\n    real_image = real_image.cuda()\n\n    if args.loss == 'wgan-gp':\n        real_predict = discriminator(real_image, step=step, alpha=alpha)\n        real_predict = real_predict.mean() - 0.001 * (real_predict ** 2).mean()\n        (-real_predict).backward()\n\n    elif args.loss == 'r1':\n        real_image.requires_grad = True\n        real_predict = discriminator(real_image, step=step, alpha=alpha)\n        real_predict = F.softplus(-real_predict).mean()\n        real_predict.backward(retain_graph=True)\n\n        grad_real = grad(\n            outputs=real_predict.sum(), inputs=real_image, create_graph=True\n        )[0]\n        grad_penalty = (\n            grad_real.view(grad_real.size(0), -1).norm(2, dim=1) ** 2\n        ).mean()\n        grad_penalty = 10 / 2 * grad_penalty\n        grad_penalty.backward()\n        grad_loss_val = grad_penalty.item()\n\n    if args.mixing and random.random() < 0.9 and False:\n        gen_in11, gen_in12, gen_in21, gen_in22 = torch.randn(\n            4, b_size, code_size, device='cuda'\n        ).chunk(4, 0)\n        gen_in1 = [gen_in11.squeeze(0), gen_in12.squeeze(0)]\n        gen_in2 = [gen_in21.squeeze(0), gen_in22.squeeze(0)]\n\n    else:\n        gen_in1, gen_in2 = torch.randn(2, b_size, code_size, device='cuda').chunk(2, 0)            \n        gen_in1 = gen_in1.squeeze(0)\n        gen_in2 = gen_in2.squeeze(0)\n\n    fake_image = generator(gen_in1, step=step, alpha=alpha)\n    fake_predict = discriminator(fake_image, step=step, alpha=alpha)\n\n    if args.loss == 'wgan-gp':\n        fake_predict = fake_predict.mean()\n        fake_predict.backward()\n\n        eps = torch.rand(b_size, 1, 1, 1).cuda()\n        x_hat = eps * real_image.data + (1 - eps) * fake_image.data\n        x_hat.requires_grad = True\n        hat_predict = discriminator(x_hat, step=step, alpha=alpha)\n        grad_x_hat = grad(\n            outputs=hat_predict.sum(), inputs=x_hat, create_graph=True\n        )[0]\n        grad_penalty = (\n            (grad_x_hat.view(grad_x_hat.size(0), -1).norm(2, dim=1) - 1) ** 2\n        ).mean()\n        grad_penalty = 10 * grad_penalty\n        grad_penalty.backward()\n        grad_loss_val = grad_penalty.item()\n        disc_loss_val = (real_predict - fake_predict).item()\n\n    elif args.loss == 'r1':\n        fake_predict = F.softplus(fake_predict).mean()\n        fake_predict.backward()\n        disc_loss_val = (real_predict + fake_predict).item()\n\n    d_optimizer.step()\n\n    if (i + 1) % n_critic == 0:\n        generator.zero_grad()\n\n        requires_grad(generator, True)\n        requires_grad(discriminator, False)\n\n        fake_image = generator(gen_in2, step=step, alpha=alpha)\n\n        predict = discriminator(fake_image, step=step, alpha=alpha)\n\n        if args.loss == 'wgan-gp':\n            loss = -predict.mean()\n\n        elif args.loss == 'r1':\n            loss = F.softplus(-predict).mean()\n\n        gen_loss_val = loss.item()\n\n        loss.backward()\n        g_optimizer.step()\n\n        requires_grad(generator, False)\n        requires_grad(discriminator, True)\n\n    if (i + 1) % 500 == 0 and Is_Show:\n        images = []\n\n        gen_i, gen_j = args.gen_sample.get(resolution, (5, 10))\n        random_z = torch.randn(gen_j, code_size).cuda()\n        with torch.no_grad():\n            for age_code in range(gen_i):\n                gen_test = torch.randn(gen_j, code_size).cuda()\n                images.append(generator(gen_test, step=step, alpha=alpha))\n       \n        gen_image_temp = vutils.make_grid(torch.cat(images, 0), nrow=10, padding=2, normalize=True)\n        gen_image_temp = gen_image_temp.cpu().numpy().transpose(1, 2, 0)\n        plt.figure(figsize=(16,6))\n        plt.imshow(gen_image_temp)\n        plt.show()\n\n    state_msg = (\n        f'Size: {8 * 2 ** step}; G: {gen_loss_val:.3f}; D: {disc_loss_val:.3f};'\n        f' Grad: {grad_loss_val:.3f}; Alpha: {alpha:.5f}'\n    )\n\n    pbar.set_description(state_msg)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from torchvision.utils import save_image\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nLOCAL = False\nif LOCAL:\n    OUTPUT_PATH = './output_images'\nelse:\n    OUTPUT_PATH = '../output_images'\nif not os.path.exists(OUTPUT_PATH):\n    os.mkdir(OUTPUT_PATH)\n    \nim_batch_size = 50\nn_images=10000\n# n_images = 5\nwith torch.no_grad():\n    for i_batch in range(0, n_images, im_batch_size):\n        gen_test = torch.randn(im_batch_size, code_size, device=device)\n    #     gen_images = (netG(gen_z)+1.)/2.\n        gen_images = (generator(gen_test, step=step, alpha=alpha) + 1.)/2.\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, :, :, :], os.path.join(OUTPUT_PATH, f'image_{i_batch+i_image:05d}.png'), normalize=True)\n\nimport shutil\nshutil.make_archive('images', 'zip', OUTPUT_PATH)","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.7.3"}},"nbformat":4,"nbformat_minor":1}