{"cells":[{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install --upgrade pytorch-lightning","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"%matplotlib inline\nimport os, glob, random\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nfrom PIL import Image\nimport shutil\n\nimport torch\nfrom torchvision.utils import make_grid\nfrom torchvision import transforms\nimport torchvision.transforms.functional as TF\nfrom torch import nn, optim\nfrom torch.optim.lr_scheduler import CosineAnnealingLR\nfrom torch.utils.data import DataLoader, Dataset\nimport pytorch_lightning as pl\nfrom torch.nn.parameter import Parameter\nfrom pytorch_lightning import Trainer\nfrom pytorch_lightning.callbacks import EarlyStopping","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"pl.__version__","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    torch.backends.cudnn.benchmark = True","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"---\n## Dataset"},{"metadata":{"trusted":true},"cell_type":"code","source":"class ImageTransform:\n    def __init__(self, img_size=256):\n        self.transform = {\n            'train': transforms.Compose([\n                transforms.Resize((img_size, img_size)),\n                transforms.ToTensor(),\n                transforms.Normalize(mean=[0.5], std=[0.5])\n            ]),\n            'test': transforms.Compose([\n                transforms.Resize((img_size, img_size)),\n                transforms.ToTensor(),\n                transforms.Normalize(mean=[0.5], std=[0.5])\n            ])}\n\n    def __call__(self, img, phase='train'):\n        img = self.transform[phase](img)\n\n        return img\n\n\n# Monet Dataset ---------------------------------------------------------------------------\nclass MonetDataset(Dataset):\n    def __init__(self, base_img_paths, style_img_paths,  transform, phase='train'):\n        self.base_img_paths = base_img_paths\n        self.style_img_paths = style_img_paths\n        self.transform = transform\n        self.phase = phase\n\n    def __len__(self):\n        return min([len(self.base_img_paths), len(self.style_img_paths)])\n\n    def __getitem__(self, idx):        \n        base_img_path = self.base_img_paths[idx]\n        style_img_path = self.style_img_paths[idx]\n        base_img = Image.open(base_img_path)\n        style_img = Image.open(style_img_path)\n\n        base_img = self.transform(base_img, self.phase)\n        style_img = self.transform(style_img, self.phase)\n\n        return base_img, style_img","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Data Module\nclass MonetDataModule(pl.LightningDataModule):\n    def __init__(self, data_dir, transform, batch_size, phase='train', seed=0, other_artist=None):\n        super(MonetDataModule, self).__init__()\n        self.data_dir = data_dir\n        self.other_artist = other_artist\n        self.transform = transform\n        self.batch_size = batch_size\n        self.phase = phase\n        self.seed = seed\n\n    def prepare_data(self):\n        self.base_img_paths = glob.glob(os.path.join(self.data_dir, 'photo_jpg', '*.jpg'))\n        if self.other_artist:\n            print(\"vangogh\")\n            self.style_img_paths = glob.glob(os.path.join(self.other_artist, 'trainA', '*.jpg'))\n        else:\n            self.style_img_paths = glob.glob(os.path.join(self.data_dir, 'monet_jpg', '*.jpg'))\n        \n    def train_dataloader(self):\n        random.seed()\n        random.shuffle(self.base_img_paths)\n        random.shuffle(self.style_img_paths)\n        random.seed(self.seed)\n        self.train_dataset = MonetDataset(self.base_img_paths, self.style_img_paths, self.transform, self.phase)\n        \n        return DataLoader(self.train_dataset,\n                          batch_size=self.batch_size,\n                          shuffle=True,\n                          pin_memory=True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Sanity Check\ndata_dir = '../input/gan-getting-started'\nother_artist = \"../input/vangogh-1\"\ntransform = ImageTransform(img_size=256)\nbatch_size = 8\n\ndm = MonetDataModule(data_dir, transform, batch_size, phase='test',other_artist=other_artist)\ndm.prepare_data()\n\ndataloader = dm.train_dataloader()\nbase, style = next(iter(dataloader))\n\nprint('Input Shape {}, {}'.format(base.size(), style.size()))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"temp = make_grid(base, nrow=4, padding=2).permute(1, 2, 0).detach().numpy()\ntemp = temp * 0.5 + 0.5\ntemp = temp * 255.0\ntemp = temp.astype(int)\n\nfig = plt.figure(figsize=(18, 8), facecolor='w')\nplt.imshow(temp)\nplt.axis('off')\nplt.title('Photo')\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"temp = make_grid(style, nrow=4, padding=2).permute(1, 2, 0).detach().numpy()\ntemp = temp * 0.5 + 0.5\ntemp = temp * 255.0\ntemp = temp.astype(int)\n\nfig = plt.figure(figsize=(18, 8), facecolor='w')\nplt.imshow(temp)\nplt.axis('off')\nplt.title('Vangogh Pictures')\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"---\n## Model"},{"metadata":{"trusted":true},"cell_type":"code","source":"class Upsample(nn.Module):\n    def __init__(self, in_channels, out_channels, kernel_size=4, stride=2, padding=1, dropout=True):\n        super(Upsample, self).__init__()\n        self.dropout = dropout\n        self.block = nn.Sequential(\n            nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride, padding, bias=nn.InstanceNorm2d),\n            nn.InstanceNorm2d(out_channels),\n            nn.ReLU(inplace=True)\n        )\n        self.dropout_layer = nn.Dropout2d(0.5)\n\n    def forward(self, x, shortcut=None):\n        x = self.block(x)\n        if self.dropout:\n            x = self.dropout_layer(x)\n\n        if shortcut is not None:\n            x = torch.cat([x, shortcut], dim=1)\n\n        return x\n\n\nclass Downsample(nn.Module):\n    def __init__(self, in_channels, out_channels, kernel_size=4, stride=2, padding=1, apply_instancenorm=True):\n        super(Downsample, self).__init__()\n        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, bias=nn.InstanceNorm2d)\n        self.norm = nn.InstanceNorm2d(out_channels)\n        self.relu = nn.LeakyReLU(0.2, inplace=True)\n        self.apply_norm = apply_instancenorm\n\n    def forward(self, x):\n        x = self.conv(x)\n        if self.apply_norm:\n            x = self.norm(x)\n        x = self.relu(x)\n\n        return x\n\n\nclass CycleGAN_Unet_Generator(nn.Module):\n    def __init__(self, filter=64):\n        super(CycleGAN_Unet_Generator, self).__init__()\n        self.downsamples = nn.ModuleList([\n            Downsample(3, filter, kernel_size=4, apply_instancenorm=False),  # (b, filter, 128, 128)\n            Downsample(filter, filter * 2),  # (b, filter * 2, 64, 64)\n            Downsample(filter * 2, filter * 4),  # (b, filter * 4, 32, 32)\n            Downsample(filter * 4, filter * 8),  # (b, filter * 8, 16, 16)\n            Downsample(filter * 8, filter * 8), # (b, filter * 8, 8, 8)\n        ])\n\n        self.upsamples = nn.ModuleList([\n            Upsample(filter * 8, filter * 8),\n            Upsample(filter * 16, filter * 4, dropout=False),\n            Upsample(filter * 8, filter * 2, dropout=False),\n            Upsample(filter * 4, filter, dropout=False)\n        ])\n\n        self.last = nn.Sequential(\n            nn.ConvTranspose2d(filter * 2, 3, kernel_size=4, stride=2, padding=1),\n            nn.Tanh()\n        )\n\n    def forward(self, x):\n        skips = []\n        for l in self.downsamples:\n            x = l(x)\n            skips.append(x)\n\n        skips = reversed(skips[:-1])\n        for l, s in zip(self.upsamples, skips):\n            x = l(x, s)\n\n        out = self.last(x)\n\n        return out\n\n\nclass CycleGAN_Discriminator(nn.Module):\n    def __init__(self, filter=64):\n        super(CycleGAN_Discriminator, self).__init__()\n\n        self.block = nn.Sequential(\n            Downsample(3, filter, kernel_size=4, stride=2, apply_instancenorm=False),\n            Downsample(filter, filter * 2, kernel_size=4, stride=2),\n            Downsample(filter * 2, filter * 4, kernel_size=4, stride=2),\n            Downsample(filter * 4, filter * 8, kernel_size=4, stride=1),\n        )\n\n        self.last = nn.Conv2d(filter * 8, 1, kernel_size=4, stride=1, padding=1)\n\n    def forward(self, x):\n        x = self.block(x)\n        x = self.last(x)\n\n        return x","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Sanity Check\nnet = CycleGAN_Unet_Generator()\n\nout = net(base)\nprint(out.size())","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Sanity Check\nnet = CycleGAN_Discriminator()\n\nout = net(base)\nprint(out.size())","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# CycleGAN - Lightning Module ---------------------------------------------------------------------------\nclass CycleGAN_LightningSystem(pl.LightningModule):\n    def __init__(self, G_basestyle, G_stylebase, D_base, D_style, lr, transform, reconstr_w=10, id_w=2):\n        super(CycleGAN_LightningSystem, self).__init__()\n        self.G_basestyle = G_basestyle\n        self.G_stylebase = G_stylebase\n        self.D_base = D_base\n        self.D_style = D_style\n        self.lr = lr\n        self.transform = transform\n        self.reconstr_w = reconstr_w\n        self.id_w = id_w\n        self.cnt_train_step = 0\n        self.step = 0\n\n        self.mae = nn.L1Loss()\n        self.generator_loss = nn.MSELoss()\n        self.discriminator_loss = nn.MSELoss()\n        self.losses = []\n        self.G_mean_losses = []\n        self.D_mean_losses = []\n        self.validity = []\n        self.reconstr = []\n        self.identity = []\n\n    def configure_optimizers(self):\n        self.g_basestyle_optimizer = optim.Adam(self.G_basestyle.parameters(), lr=self.lr['G'], betas=(0.5, 0.999))\n        self.g_stylebase_optimizer = optim.Adam(self.G_stylebase.parameters(), lr=self.lr['G'], betas=(0.5, 0.999))\n        self.d_base_optimizer = optim.Adam(self.D_base.parameters(), lr=self.lr['D'], betas=(0.5, 0.999))\n        self.d_style_optimizer = optim.Adam(self.D_style.parameters(), lr=self.lr['D'], betas=(0.5, 0.999))\n\n        return [self.g_basestyle_optimizer, self.g_stylebase_optimizer, self.d_base_optimizer, self.d_style_optimizer], []\n\n    def training_step(self, batch, batch_idx, optimizer_idx):\n        base_img, style_img = batch\n        b = base_img.size()[0]\n\n        valid = torch.ones(b, 1, 30, 30).cuda()\n        fake = torch.zeros(b, 1, 30, 30).cuda()\n\n        # Train Generator\n        if optimizer_idx == 0 or optimizer_idx == 1:\n            # Validity\n            # MSELoss\n            val_base = self.generator_loss(self.D_base(self.G_stylebase(style_img)), valid)\n            val_style = self.generator_loss(self.D_style(self.G_basestyle(base_img)), valid)\n            val_loss = (val_base + val_style) / 2\n\n            # Reconstruction\n            reconstr_base = self.mae(self.G_stylebase(self.G_basestyle(base_img)), base_img)\n            reconstr_style = self.mae(self.G_basestyle(self.G_stylebase(style_img)), style_img)\n            reconstr_loss = (reconstr_base + reconstr_style) / 2\n\n            # Identity\n            id_base = self.mae(self.G_stylebase(base_img), base_img)\n            id_style = self.mae(self.G_basestyle(style_img), style_img)\n            id_loss = (id_base + id_style) / 2\n\n            # Loss Weight\n            G_loss = val_loss + self.reconstr_w * reconstr_loss + self.id_w * id_loss\n\n            return {'loss': G_loss, 'validity': val_loss, 'reconstr': reconstr_loss, 'identity': id_loss}\n\n        # Train Discriminator\n        elif optimizer_idx == 2 or optimizer_idx == 3:\n            # MSELoss\n            D_base_gen_loss = self.discriminator_loss(self.D_base(self.G_stylebase(style_img)), fake)\n            D_style_gen_loss = self.discriminator_loss(self.D_style(self.G_basestyle(base_img)), fake)\n            D_base_valid_loss = self.discriminator_loss(self.D_base(base_img), valid)\n            D_style_valid_loss = self.discriminator_loss(self.D_style(style_img), valid)\n            \n            D_gen_loss = (D_base_gen_loss + D_style_gen_loss) / 2\n            \n            # Loss Weight\n            D_loss = (D_gen_loss + D_base_valid_loss + D_style_valid_loss) / 3\n\n            # Count up\n            self.cnt_train_step += 1\n\n            return {'loss': D_loss}\n\n    def training_epoch_end(self, outputs):\n        self.step += 1\n        \n        avg_loss = sum([torch.stack([x['loss'] for x in outputs[i]]).mean().item() / 4 for i in range(4)])\n        G_mean_loss = sum([torch.stack([x['loss'] for x in outputs[i]]).mean().item() / 2 for i in [0, 1]])\n        D_mean_loss = sum([torch.stack([x['loss'] for x in outputs[i]]).mean().item() / 2 for i in [2, 3]])\n        validity = sum([torch.stack([x['validity'] for x in outputs[i]]).mean().item() / 2 for i in [0, 1]])\n        reconstr = sum([torch.stack([x['reconstr'] for x in outputs[i]]).mean().item() / 2 for i in [0, 1]])\n        identity = sum([torch.stack([x['identity'] for x in outputs[i]]).mean().item() / 2 for i in [0, 1]])\n            \n        self.losses.append(avg_loss)\n        self.G_mean_losses.append(G_mean_loss)\n        self.D_mean_losses.append(D_mean_loss)\n        self.validity.append(validity)\n        self.reconstr.append(reconstr)\n        self.identity.append(identity)\n        self.log('validity_loss', validity)\n        self.log('reconstr_loss', reconstr)\n        self.log('identity_loss', identity)\n        self.log('D_mean_loss', D_mean_loss)\n        self.log('G_mean_loss', G_mean_loss)\n        self.log('avg_loss', avg_loss)\n        \n        if self.step % 10 == 0:\n            # Display Model Output\n            target_img_paths = glob.glob('../input/gan-getting-started/photo_jpg/*.jpg')[:4]\n            target_imgs = [self.transform(Image.open(path), phase='test') for path in target_img_paths]\n            target_imgs = torch.stack(target_imgs, dim=0)\n            target_imgs = target_imgs.cuda()\n\n            gen_imgs = self.G_basestyle(target_imgs)\n            gen_img = torch.cat([target_imgs, gen_imgs], dim=0)\n\n            # Reverse Normalization\n            gen_img = gen_img * 0.5 + 0.5\n            gen_img = gen_img * 255\n\n            joined_images_tensor = make_grid(gen_img, nrow=4, padding=2)\n\n            joined_images = joined_images_tensor.detach().cpu().numpy().astype(int)\n            joined_images = np.transpose(joined_images, [1,2,0])\n\n            # Visualize\n            fig = plt.figure(figsize=(18, 8))\n            plt.imshow(joined_images)\n            plt.axis('off')\n            plt.title(f'Epoch {self.step}')\n            plt.show()\n            plt.clf()\n            plt.close()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def init_weights(net, init_type='normal', init_gain=0.02):\n    \"\"\"Initialize network weights.\n    Parameters:\n        net (network)   -- network to be initialized\n        init_type (str) -- the name of an initialization method: normal | xavier | kaiming | orthogonal\n        init_gain (float)    -- scaling factor for normal, xavier and orthogonal.\n    We use 'normal' in the original pix2pix and CycleGAN paper. But xavier and kaiming might\n    work better for some applications. Feel free to try yourself.\n    \"\"\"\n    def init_func(m):  # define the initialization function\n        classname = m.__class__.__name__\n        if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1):\n            if init_type == 'normal':\n                nn.init.normal_(m.weight.data, 0.0, init_gain)\n            elif init_type == 'xavier':\n                nn.init.xavier_normal_(m.weight.data, gain=init_gain)\n            elif init_type == 'kaiming':\n                nn.init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')\n            elif init_type == 'orthogonal':\n                nn.init.orthogonal_(m.weight.data, gain=init_gain)\n            else:\n                raise NotImplementedError('initialization method [%s] is not implemented' % init_type)\n            if hasattr(m, 'bias') and m.bias is not None:\n                nn.init.constant_(m.bias.data, 0.0)\n        elif classname.find('BatchNorm2d') != -1:  # BatchNorm Layer's weight is not a matrix; only normal distribution applies.\n            nn.init.normal_(m.weight.data, 1.0, init_gain)\n            nn.init.constant_(m.bias.data, 0.0)\n\n    net.apply(init_func)  # apply the initialization function <init_func>","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class ResnetGenerator(nn.Module):\n    def __init__(self, input_nc, output_nc, ngf=64, n_blocks=6, img_size=256, light=False):\n        assert(n_blocks >= 0)\n        super(ResnetGenerator, self).__init__()\n        self.input_nc = input_nc\n        self.output_nc = output_nc\n        self.ngf = ngf\n        self.n_blocks = n_blocks\n        self.img_size = img_size\n        self.light = light\n\n        DownBlock = []\n        DownBlock += [nn.ReflectionPad2d(3),\n                      nn.Conv2d(input_nc, ngf, kernel_size=7, stride=1, padding=0, bias=False),\n                      nn.InstanceNorm2d(ngf),\n                      nn.ReLU(True)]\n\n        # Down-Sampling\n        n_downsampling = 2\n        for i in range(n_downsampling):\n            mult = 2**i\n            DownBlock += [nn.ReflectionPad2d(1),\n                          nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=0, bias=False),\n                          nn.InstanceNorm2d(ngf * mult * 2),\n                          nn.ReLU(True)]\n\n        # Down-Sampling Bottleneck\n        mult = 2**n_downsampling\n        for i in range(n_blocks):\n            DownBlock += [ResnetBlock(ngf * mult, use_bias=False)]\n\n        # Class Activation Map\n        self.gap_fc = nn.Linear(ngf * mult, 1, bias=False)\n        self.gmp_fc = nn.Linear(ngf * mult, 1, bias=False)\n        self.conv1x1 = nn.Conv2d(ngf * mult * 2, ngf * mult, kernel_size=1, stride=1, bias=True)\n        self.relu = nn.ReLU(True)\n\n        # Gamma, Beta block\n        if self.light:\n            FC = [nn.Linear(ngf * mult, ngf * mult, bias=False),\n                  nn.ReLU(True),\n                  nn.Linear(ngf * mult, ngf * mult, bias=False),\n                  nn.ReLU(True)]\n        else:\n            FC = [nn.Linear(img_size // mult * img_size // mult * ngf * mult, ngf * mult, bias=False),\n                  nn.ReLU(True),\n                  nn.Linear(ngf * mult, ngf * mult, bias=False),\n                  nn.ReLU(True)]\n        self.gamma = nn.Linear(ngf * mult, ngf * mult, bias=False)\n        self.beta = nn.Linear(ngf * mult, ngf * mult, bias=False)\n\n        # Up-Sampling Bottleneck\n        for i in range(n_blocks):\n            setattr(self, 'UpBlock1_' + str(i+1), ResnetAdaILNBlock(ngf * mult, use_bias=False))\n\n        # Up-Sampling\n        UpBlock2 = []\n        for i in range(n_downsampling):\n            mult = 2**(n_downsampling - i)\n            UpBlock2 += [nn.Upsample(scale_factor=2, mode='nearest'),\n                         nn.ReflectionPad2d(1),\n                         nn.Conv2d(ngf * mult, int(ngf * mult / 2), kernel_size=3, stride=1, padding=0, bias=False),\n                         ILN(int(ngf * mult / 2)),\n                         nn.ReLU(True)]\n\n        UpBlock2 += [nn.ReflectionPad2d(3),\n                     nn.Conv2d(ngf, output_nc, kernel_size=7, stride=1, padding=0, bias=False),\n                     nn.Tanh()]\n\n        self.DownBlock = nn.Sequential(*DownBlock)\n        self.FC = nn.Sequential(*FC)\n        self.UpBlock2 = nn.Sequential(*UpBlock2)\n\n    def forward(self, input):\n        x = self.DownBlock(input)\n\n        gap = torch.nn.functional.adaptive_avg_pool2d(x, 1)\n        gap_logit = self.gap_fc(gap.view(x.shape[0], -1))\n        gap_weight = list(self.gap_fc.parameters())[0]\n        gap = x * gap_weight.unsqueeze(2).unsqueeze(3)\n\n        gmp = torch.nn.functional.adaptive_max_pool2d(x, 1)\n        gmp_logit = self.gmp_fc(gmp.view(x.shape[0], -1))\n        gmp_weight = list(self.gmp_fc.parameters())[0]\n        gmp = x * gmp_weight.unsqueeze(2).unsqueeze(3)\n\n        cam_logit = torch.cat([gap_logit, gmp_logit], 1)\n        x = torch.cat([gap, gmp], 1)\n        x = self.relu(self.conv1x1(x))\n\n        heatmap = torch.sum(x, dim=1, keepdim=True)\n\n        if self.light:\n            x_ = torch.nn.functional.adaptive_avg_pool2d(x, 1)\n            x_ = self.FC(x_.view(x_.shape[0], -1))\n        else:\n            x_ = self.FC(x.view(x.shape[0], -1))\n        gamma, beta = self.gamma(x_), self.beta(x_)\n\n\n        for i in range(self.n_blocks):\n            x = getattr(self, 'UpBlock1_' + str(i+1))(x, gamma, beta)\n        out = self.UpBlock2(x)\n\n        return out#, cam_logit, heatmap\n\n\nclass ResnetBlock(nn.Module):\n    def __init__(self, dim, use_bias):\n        super(ResnetBlock, self).__init__()\n        conv_block = []\n        conv_block += [nn.ReflectionPad2d(1),\n                       nn.Conv2d(dim, dim, kernel_size=3, stride=1, padding=0, bias=use_bias),\n                       nn.InstanceNorm2d(dim),\n                       nn.ReLU(True)]\n\n        conv_block += [nn.ReflectionPad2d(1),\n                       nn.Conv2d(dim, dim, kernel_size=3, stride=1, padding=0, bias=use_bias),\n                       nn.InstanceNorm2d(dim)]\n\n        self.conv_block = nn.Sequential(*conv_block)\n\n    def forward(self, x):\n        out = x + self.conv_block(x)\n        return out\n\n\nclass ResnetAdaILNBlock(nn.Module):\n    def __init__(self, dim, use_bias):\n        super(ResnetAdaILNBlock, self).__init__()\n        self.pad1 = nn.ReflectionPad2d(1)\n        self.conv1 = nn.Conv2d(dim, dim, kernel_size=3, stride=1, padding=0, bias=use_bias)\n        self.norm1 = adaILN(dim)\n        self.relu1 = nn.ReLU(True)\n\n        self.pad2 = nn.ReflectionPad2d(1)\n        self.conv2 = nn.Conv2d(dim, dim, kernel_size=3, stride=1, padding=0, bias=use_bias)\n        self.norm2 = adaILN(dim)\n\n    def forward(self, x, gamma, beta):\n        out = self.pad1(x)\n        out = self.conv1(out)\n        out = self.norm1(out, gamma, beta)\n        out = self.relu1(out)\n        out = self.pad2(out)\n        out = self.conv2(out)\n        out = self.norm2(out, gamma, beta)\n\n        return out + x\n\n    \nclass adaILN(nn.Module):\n    def __init__(self, num_features, eps=1e-5):\n        super(adaILN, self).__init__()\n        self.eps = eps\n        self.rho = Parameter(torch.Tensor(1, num_features, 1, 1))\n        self.rho.data.fill_(0.9)\n\n    def forward(self, input, gamma, beta):\n        in_mean, in_var = torch.mean(input, dim=[2, 3], keepdim=True), torch.var(input, dim=[2, 3], keepdim=True)\n        out_in = (input - in_mean) / torch.sqrt(in_var + self.eps)\n        ln_mean, ln_var = torch.mean(input, dim=[1, 2, 3], keepdim=True), torch.var(input, dim=[1, 2, 3], keepdim=True)\n        out_ln = (input - ln_mean) / torch.sqrt(ln_var + self.eps)\n        out = self.rho.expand(input.shape[0], -1, -1, -1) * out_in + (1-self.rho.expand(input.shape[0], -1, -1, -1)) * out_ln\n        out = out * gamma.unsqueeze(2).unsqueeze(3) + beta.unsqueeze(2).unsqueeze(3)\n\n        return out\n    \nclass ILN(nn.Module):\n    def __init__(self, num_features, eps=1e-5):\n        super(ILN, self).__init__()\n        self.eps = eps\n        self.rho = Parameter(torch.Tensor(1, num_features, 1, 1))\n        self.gamma = Parameter(torch.Tensor(1, num_features, 1, 1))\n        self.beta = Parameter(torch.Tensor(1, num_features, 1, 1))\n        self.rho.data.fill_(0.0)\n        self.gamma.data.fill_(1.0)\n        self.beta.data.fill_(0.0)\n\n    def forward(self, input):\n        in_mean, in_var = torch.mean(input, dim=[2, 3], keepdim=True), torch.var(input, dim=[2, 3], keepdim=True)\n        out_in = (input - in_mean) / torch.sqrt(in_var + self.eps)\n        ln_mean, ln_var = torch.mean(input, dim=[1, 2, 3], keepdim=True), torch.var(input, dim=[1, 2, 3], keepdim=True)\n        out_ln = (input - ln_mean) / torch.sqrt(ln_var + self.eps)\n        out = self.rho.expand(input.shape[0], -1, -1, -1) * out_in + (1-self.rho.expand(input.shape[0], -1, -1, -1)) * out_ln\n        out = out * self.gamma.expand(input.shape[0], -1, -1, -1) + self.beta.expand(input.shape[0], -1, -1, -1)\n\n        return out\n\n\nclass Discriminator(nn.Module):\n    def __init__(self, input_nc, ndf=64, n_layers=5):\n        super(Discriminator, self).__init__()\n        model = [nn.ReflectionPad2d(1),\n                 nn.utils.spectral_norm(\n                 nn.Conv2d(input_nc, ndf, kernel_size=4, stride=2, padding=0, bias=True)),\n                 nn.LeakyReLU(0.2, True)]\n\n        for i in range(1, n_layers - 2):\n            mult = 2 ** (i - 1)\n            model += [nn.ReflectionPad2d(1),\n                      nn.utils.spectral_norm(\n                      nn.Conv2d(ndf * mult, ndf * mult * 2, kernel_size=4, stride=2, padding=0, bias=True)),\n                      nn.LeakyReLU(0.2, True)]\n\n        mult = 2 ** (n_layers - 2 - 1)\n        model += [nn.ReflectionPad2d(1),\n                  nn.utils.spectral_norm(\n                  nn.Conv2d(ndf * mult, ndf * mult * 2, kernel_size=4, stride=1, padding=0, bias=True)),\n                  nn.LeakyReLU(0.2, True)]\n\n        # Class Activation Map\n        mult = 2 ** (n_layers - 2)\n        self.gap_fc = nn.utils.spectral_norm(nn.Linear(ndf * mult, 1, bias=False))\n        self.gmp_fc = nn.utils.spectral_norm(nn.Linear(ndf * mult, 1, bias=False))\n        self.conv1x1 = nn.Conv2d(ndf * mult * 2, ndf * mult, kernel_size=1, stride=1, bias=True)\n        self.leaky_relu = nn.LeakyReLU(0.2, True)\n\n        self.pad = nn.ReflectionPad2d(1)\n        self.conv = nn.utils.spectral_norm(\n            nn.Conv2d(ndf * mult, 1, kernel_size=4, stride=1, padding=0, bias=False))\n\n        self.model = nn.Sequential(*model)\n\n    def forward(self, input):\n        x = self.model(input)\n\n        gap = torch.nn.functional.adaptive_avg_pool2d(x, 1)\n        gap_logit = self.gap_fc(gap.view(x.shape[0], -1))\n        gap_weight = list(self.gap_fc.parameters())[0]\n        gap = x * gap_weight.unsqueeze(2).unsqueeze(3)\n\n        gmp = torch.nn.functional.adaptive_max_pool2d(x, 1)\n        gmp_logit = self.gmp_fc(gmp.view(x.shape[0], -1))\n        gmp_weight = list(self.gmp_fc.parameters())[0]\n        gmp = x * gmp_weight.unsqueeze(2).unsqueeze(3)\n\n        cam_logit = torch.cat([gap_logit, gmp_logit], 1)\n        x = torch.cat([gap, gmp], 1)\n        x = self.leaky_relu(self.conv1x1(x))\n\n        heatmap = torch.sum(x, dim=1, keepdim=True)\n\n        x = self.pad(x)\n        out = self.conv(x)\n\n        return out, cam_logit, heatmap","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"---\n## Train"},{"metadata":{"trusted":true},"cell_type":"code","source":"# Config  -----------------------------------------------------------------\ndata_dir = '../input/gan-getting-started'\nother_artist = \"../input/vangogh-1\"\ntransform = ImageTransform(img_size=256)\nbatch_size = 1\nlr = {\n    'G': 0.0002,\n    'D': 0.0002\n}\nepoch = 40\nseed = 42\nreconstr_w = 10\nid_w = 5\nseed_everything(seed)\n\n# DataModule  -----------------------------------------------------------------\ndm = MonetDataModule(data_dir, transform, batch_size, seed=seed, other_artist=other_artist)\n\n# G_basestyle = CycleGAN_Unet_Generator()\n# G_stylebase = CycleGAN_Unet_Generator()\n# D_base = CycleGAN_Discriminator()\n# D_style = CycleGAN_Discriminator()\n\nG_basestyle = ResnetGenerator(input_nc=3, output_nc=3, ngf=64, n_blocks=4, img_size=256, light=True)\nG_stylebase = ResnetGenerator(input_nc=3, output_nc=3, ngf=64, n_blocks=4, img_size=256, light=True)\nD_base = CycleGAN_Discriminator()\nD_style = CycleGAN_Discriminator()\n\n\n# Init Weight  --------------------------------------------------------------\nfor net in [G_basestyle, G_stylebase, D_base, D_style]:\n    init_weights(net, init_type='normal')\n\n# LightningModule  --------------------------------------------------------------\nmodel = CycleGAN_LightningSystem(G_basestyle, G_stylebase, D_base, D_style, \n                                 lr, transform, reconstr_w, id_w)\n\nearly_stop_callback = EarlyStopping(\n  monitor='G_mean_loss',\n  min_delta=0.00,\n  patience=3,\n  verbose=False,\n  mode='min'\n)\n\n\n# Trainer  --------------------------------------------------------------\ntrainer = Trainer(\n    logger=True,\n    max_epochs=epoch,\n    gpus=1,\n    checkpoint_callback=False,\n    callbacks=[early_stop_callback],\n    reload_dataloaders_every_epoch=True,\n    num_sanity_val_steps=0,  # Skip Sanity Check\n)\n\n\n# Train\ntrainer.fit(model, datamodule=dm)\ntrainer.save_checkpoint(\"/kaggle/working/vangogh_1.ckpt\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"PATH = \"/kaggle/working/vangogh_1.ckpt\"\n\nnew_model = CycleGAN_LightningSystem.load_from_checkpoint(checkpoint_path=PATH, G_basestyle=G_basestyle, G_stylebase=G_stylebase, \n                                                          D_base=D_base, D_style=D_style, lr=lr, transform=transform)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"dm = MonetDataModule(data_dir, transform, batch_size, seed=seed)\n# trainer.max_epochs = 160\n\nearly_stop_callback = EarlyStopping(\n  monitor='G_mean_loss',\n  min_delta=0.00,\n  patience=3,\n  verbose=False,\n  mode='min'\n)\n\nmonet_trainer = Trainer(\n    logger=True,\n    max_epochs=75,\n    gpus=1,\n    checkpoint_callback=False,\n    callbacks=[early_stop_callback],\n    reload_dataloaders_every_epoch=True,\n    num_sanity_val_steps=0,  # Skip Sanity Check\n)\nmonet_trainer.fit(new_model, datamodule=dm)\nmonet_trainer.save_checkpoint(\"/kaggle/working/monet.ckpt\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def submit(model, transform):\n    os.makedirs('../images', exist_ok=True)\n    net = model.G_basestyle\n    \n    net.eval()\n    photo_img_paths = glob.glob('../input/gan-getting-started/photo_jpg/*.jpg')\n    \n    for path in photo_img_paths:\n        photo_id = path.split('/')[-1]\n        img = transform(Image.open(path), phase='test')\n        \n        gen_img = net(img.unsqueeze(0))[0]\n        \n        # Reverse Normalization\n        gen_img = gen_img * 0.5 + 0.5\n        gen_img = gen_img * 255\n        gen_img = gen_img.detach().cpu().numpy().astype(np.uint8)\n        \n        gen_img = np.transpose(gen_img, [1,2,0])\n        \n        gen_img = Image.fromarray(gen_img)\n        gen_img.save(os.path.join('../images', photo_id))\n        \n    # Make Zipfile\n    shutil.make_archive(\"/kaggle/working/images\", 'zip', \"/kaggle/images\")\n    \n    # Delete Origin file\n    shutil.rmtree('../images')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"submit(model, transform)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Loss Plot\nfig, axes = plt.subplots(ncols=1, nrows=2, figsize=(18, 12), facecolor='w')\nepoch_num = len(model.losses)\n\naxes[0].plot(np.arange(epoch_num), model.losses, label='overall')\naxes[0].plot(np.arange(epoch_num), model.G_mean_losses, label='generator')\naxes[0].plot(np.arange(epoch_num), model.D_mean_losses, label='discriminator')\naxes[0].legend()\naxes[0].set_xlabel('Epoch')\n\naxes[1].plot(np.arange(epoch_num), model.validity, label='validity')\naxes[1].plot(np.arange(epoch_num), model.reconstr, label='reconstr')\naxes[1].plot(np.arange(epoch_num), model.identity, label='identity')\naxes[1].legend()\naxes[1].set_xlabel('Epoch')\n\nplt.show()","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}