{"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_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install pytorch-lightning-bolts -q","metadata":{"execution":{"iopub.status.busy":"2021-06-28T16:09:08.039897Z","iopub.execute_input":"2021-06-28T16:09:08.040542Z","iopub.status.idle":"2021-06-28T16:09:13.926327Z","shell.execute_reply.started":"2021-06-28T16:09:08.040498Z","shell.execute_reply":"2021-06-28T16:09:13.925231Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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 pytorch_lightning import Trainer\nfrom pytorch_lightning.callbacks import EarlyStopping\n\nfrom pl_bolts.models.self_supervised import SimSiam \nfrom pl_bolts.models.self_supervised.simclr import SimCLREvalDataTransform, SimCLRTrainDataTransform","metadata":{"_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","papermill":{"duration":2.696591,"end_time":"2020-12-27T17:10:28.381186","exception":false,"start_time":"2020-12-27T17:10:25.684595","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2021-06-28T16:21:46.717318Z","iopub.execute_input":"2021-06-28T16:21:46.71764Z","iopub.status.idle":"2021-06-28T16:21:46.729001Z","shell.execute_reply.started":"2021-06-28T16:21:46.71761Z","shell.execute_reply":"2021-06-28T16:21:46.728165Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"pl.__version__","metadata":{"papermill":{"duration":0.028728,"end_time":"2020-12-27T17:10:28.425665","exception":false,"start_time":"2020-12-27T17:10:28.396937","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2021-06-28T16:09:14.908243Z","iopub.execute_input":"2021-06-28T16:09:14.908576Z","iopub.status.idle":"2021-06-28T16:09:14.916985Z","shell.execute_reply.started":"2021-06-28T16:09:14.90854Z","shell.execute_reply":"2021-06-28T16:09:14.916165Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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","metadata":{"papermill":{"duration":0.026511,"end_time":"2020-12-27T17:10:28.467899","exception":false,"start_time":"2020-12-27T17:10:28.441388","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2021-06-28T16:09:14.920009Z","iopub.execute_input":"2021-06-28T16:09:14.920301Z","iopub.status.idle":"2021-06-28T16:09:14.926016Z","shell.execute_reply.started":"2021-06-28T16:09:14.920275Z","shell.execute_reply":"2021-06-28T16:09:14.925107Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"---\n## Dataset","metadata":{"papermill":{"duration":0.015273,"end_time":"2020-12-27T17:10:28.498557","exception":false,"start_time":"2020-12-27T17:10:28.483284","status":"completed"},"tags":[]}},{"cell_type":"code","source":"class ImageTransform:\n    def __init__(self, img_size=256):\n        self.transform = {\n            'ssl': transforms.Compose([\n                transforms.Resize((img_size, img_size)),\n                transforms.RandomHorizontalFlip(p=0.5),\n                transforms.ToTensor(),\n                transforms.Normalize(mean=[0.5], std=[0.5])\n            ]),\n            'train': transforms.Compose([\n                transforms.Resize((img_size, img_size)),\n                transforms.RandomHorizontalFlip(p=0.5),\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        if self.phase==\"ssl\":\n                return base_img, style_img , 1.0\n        return base_img, style_img","metadata":{"papermill":{"duration":0.031267,"end_time":"2020-12-27T17:10:28.544584","exception":false,"start_time":"2020-12-27T17:10:28.513317","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2021-06-28T16:18:33.787549Z","iopub.execute_input":"2021-06-28T16:18:33.7879Z","iopub.status.idle":"2021-06-28T16:18:33.799002Z","shell.execute_reply.started":"2021-06-28T16:18:33.787863Z","shell.execute_reply":"2021-06-28T16:18:33.797895Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Data Module\nclass MonetDataModule(pl.LightningDataModule):\n    def __init__(self, data_dir, transform, batch_size, monet_limit=0, \n                 base_img_paths=None, style_img_paths=None, phase='train', seed=0):\n        super(MonetDataModule, self).__init__()\n        self.data_dir = data_dir\n        self.transform = transform\n        self.batch_size = batch_size\n        self.phase = phase\n        self.seed = seed\n        self.monet_limit = monet_limit\n        \n        if base_img_paths is None and style_img_paths is None:\n            self.base_img_paths = glob.glob(os.path.join(self.data_dir, 'photo_jpg', '*.jpg'))\n            self.style_img_paths = glob.glob(os.path.join(self.data_dir, 'monet_jpg', '*.jpg'))\n\n            random.seed()\n            random.shuffle(self.base_img_paths)\n            random.shuffle(self.style_img_paths)\n            random.seed(self.seed)\n\n            if 0 < self.monet_limit:\n                self.style_img_paths = self.style_img_paths[:self.monet_limit]\n        else:\n            self.base_img_paths = base_img_paths\n            self.style_img_paths = style_img_paths\n            \n    def train_dataloader(self):\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\n                         )","metadata":{"papermill":{"duration":0.028239,"end_time":"2020-12-27T17:10:28.587392","exception":false,"start_time":"2020-12-27T17:10:28.559153","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2021-06-28T16:18:35.838292Z","iopub.execute_input":"2021-06-28T16:18:35.838646Z","iopub.status.idle":"2021-06-28T16:18:35.849103Z","shell.execute_reply.started":"2021-06-28T16:18:35.838613Z","shell.execute_reply":"2021-06-28T16:18:35.847974Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Sanity Check\ndata_dir = '../input/gan-getting-started'\ntransform = ImageTransform(img_size=256)\nbatch_size = 8\n\ndm = MonetDataModule(data_dir, transform, batch_size, phase='test')\n\ndataloader = dm.train_dataloader()\nbase, style = next(iter(dataloader))\n\nprint('Input Shape {}, {}'.format(base.size(), style.size()))","metadata":{"papermill":{"duration":5.1579,"end_time":"2020-12-27T17:10:33.760319","exception":false,"start_time":"2020-12-27T17:10:28.602419","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2021-06-28T16:18:36.080301Z","iopub.execute_input":"2021-06-28T16:18:36.080554Z","iopub.status.idle":"2021-06-28T16:18:36.184357Z","shell.execute_reply.started":"2021-06-28T16:18:36.08053Z","shell.execute_reply":"2021-06-28T16:18:36.183464Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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()","metadata":{"papermill":{"duration":0.475999,"end_time":"2020-12-27T17:10:34.253602","exception":false,"start_time":"2020-12-27T17:10:33.777603","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2021-06-28T16:18:36.2554Z","iopub.execute_input":"2021-06-28T16:18:36.255684Z","iopub.status.idle":"2021-06-28T16:18:36.615294Z","shell.execute_reply.started":"2021-06-28T16:18:36.255654Z","shell.execute_reply":"2021-06-28T16:18:36.614384Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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('Monet Pictures')\nplt.show()","metadata":{"papermill":{"duration":0.411615,"end_time":"2020-12-27T17:10:34.697398","exception":false,"start_time":"2020-12-27T17:10:34.285783","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2021-06-28T16:18:36.616828Z","iopub.execute_input":"2021-06-28T16:18:36.617323Z","iopub.status.idle":"2021-06-28T16:18:36.99391Z","shell.execute_reply.started":"2021-06-28T16:18:36.617287Z","shell.execute_reply":"2021-06-28T16:18:36.992936Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"---\n## Model","metadata":{"papermill":{"duration":0.05777,"end_time":"2020-12-27T17:10:34.809284","exception":false,"start_time":"2020-12-27T17:10:34.751514","status":"completed"},"tags":[]}},{"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","metadata":{"papermill":{"duration":0.088394,"end_time":"2020-12-27T17:10:34.95324","exception":false,"start_time":"2020-12-27T17:10:34.864846","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2021-06-28T16:18:36.995579Z","iopub.execute_input":"2021-06-28T16:18:36.995927Z","iopub.status.idle":"2021-06-28T16:18:37.0224Z","shell.execute_reply.started":"2021-06-28T16:18:36.995893Z","shell.execute_reply":"2021-06-28T16:18:37.021459Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Sanity Check\nnet = CycleGAN_Unet_Generator()\n\nout = net(base)\nprint(out.size())","metadata":{"papermill":{"duration":1.733534,"end_time":"2020-12-27T17:10:36.738819","exception":false,"start_time":"2020-12-27T17:10:35.005285","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2021-06-28T16:18:37.307972Z","iopub.execute_input":"2021-06-28T16:18:37.308253Z","iopub.status.idle":"2021-06-28T16:18:38.780877Z","shell.execute_reply.started":"2021-06-28T16:18:37.308227Z","shell.execute_reply":"2021-06-28T16:18:38.780029Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Sanity Check\nnet = CycleGAN_Discriminator()\n\nout = net(base)\nprint(out.size())","metadata":{"papermill":{"duration":0.781824,"end_time":"2020-12-27T17:10:37.574216","exception":false,"start_time":"2020-12-27T17:10:36.792392","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2021-06-28T16:18:38.785355Z","iopub.execute_input":"2021-06-28T16:18:38.787488Z","iopub.status.idle":"2021-06-28T16:18:39.572316Z","shell.execute_reply.started":"2021-06-28T16:18:38.787446Z","shell.execute_reply":"2021-06-28T16:18:39.57139Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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        \n        if True:\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()\n\n        return None","metadata":{"papermill":{"duration":0.104365,"end_time":"2020-12-27T17:10:37.730969","exception":false,"start_time":"2020-12-27T17:10:37.626604","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2021-06-28T16:18:39.575037Z","iopub.execute_input":"2021-06-28T16:18:39.575502Z","iopub.status.idle":"2021-06-28T16:18:39.604981Z","shell.execute_reply.started":"2021-06-28T16:18:39.575461Z","shell.execute_reply":"2021-06-28T16:18:39.603675Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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>","metadata":{"papermill":{"duration":0.073807,"end_time":"2020-12-27T17:10:37.859578","exception":false,"start_time":"2020-12-27T17:10:37.785771","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2021-06-28T16:18:39.606606Z","iopub.execute_input":"2021-06-28T16:18:39.607139Z","iopub.status.idle":"2021-06-28T16:18:39.620632Z","shell.execute_reply.started":"2021-06-28T16:18:39.607029Z","shell.execute_reply":"2021-06-28T16:18:39.619545Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"___\n## Config","metadata":{}},{"cell_type":"code","source":"# Config  -----------------------------------------------------------------\ndata_dir = '../input/gan-getting-started'\nbatch_size = 16\nlr = {\n    'G': 0.0002,\n    'D': 0.0002\n}\nepoch = 1000\nseed = 42\nreconstr_w = 10\nid_w = 5\nseed_everything(seed)\n\nmonet_limits = list(range(0,35,5))[1:]\nassert max(monet_limits)<=30\n\n\ntransform = ImageTransform(img_size=256)\n\n# DataModule  -----------------------------------------------------------------\ndm = MonetDataModule(data_dir, transform, batch_size, monet_limit=monet_limits[-1], seed=seed, phase=\"ssl\")\nbase_img_paths, style_img_paths = dm.base_img_paths, dm.style_img_paths\nassert len(style_img_paths) == monet_limits[-1]","metadata":{"execution":{"iopub.status.busy":"2021-06-28T16:18:39.797977Z","iopub.execute_input":"2021-06-28T16:18:39.798235Z","iopub.status.idle":"2021-06-28T16:18:39.839439Z","shell.execute_reply.started":"2021-06-28T16:18:39.798211Z","shell.execute_reply":"2021-06-28T16:18:39.838611Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"G_basestyle = CycleGAN_Unet_Generator()\nG_stylebase = CycleGAN_Unet_Generator()\nD_base = CycleGAN_Discriminator()\nD_style = CycleGAN_Discriminator()\n\n# Init Weight  --------------------------------------------------------------\nfor net in [G_basestyle, G_stylebase, D_base, D_style]:\n    init_weights(net, init_type='normal')","metadata":{"execution":{"iopub.status.busy":"2021-06-28T16:18:40.545812Z","iopub.execute_input":"2021-06-28T16:18:40.546119Z","iopub.status.idle":"2021-06-28T16:18:41.143607Z","shell.execute_reply.started":"2021-06-28T16:18:40.546092Z","shell.execute_reply":"2021-06-28T16:18:41.142702Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"---\n## Train SimSiam","metadata":{}},{"cell_type":"code","source":"# # Sanity Check\n# net = CycleGAN_Unet_Generator()\n# downsamples = nn.Sequential(*list(net.downsamples.children()))\n# out = downsamples(base)\n# print(out.size())","metadata":{"execution":{"iopub.status.busy":"2021-06-28T16:18:41.667615Z","iopub.execute_input":"2021-06-28T16:18:41.667934Z","iopub.status.idle":"2021-06-28T16:18:42.140201Z","shell.execute_reply.started":"2021-06-28T16:18:41.667904Z","shell.execute_reply":"2021-06-28T16:18:42.139193Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# model = SimSiam(arch=\"resnet18\",gpus=1,num_samples=30,batch_size=batch_size,dataset=\"imagenet2012\")\n# net = CycleGAN_Unet_Generator()\n# downsamples = nn.Sequential(*list(net.downsamples.children()))\n# model.encoder = downsamples\n# dm.train_transforms = SimCLRTrainDataTransform(256)\n# dm.val_transforms = SimCLREvalDataTransform(256)\n\n# trainer = pl.Trainer(\n#     logger=False,\n#     max_epochs=1000,\n#     gpus=1,\n#     checkpoint_callback=False,\n#     reload_dataloaders_every_epoch=True,\n#     num_sanity_val_steps=0,  # Skip Sanity Check\n# )\n# trainer.fit(model, datamodule=dm)","metadata":{"execution":{"iopub.status.busy":"2021-06-28T16:21:59.976057Z","iopub.execute_input":"2021-06-28T16:21:59.976372Z","iopub.status.idle":"2021-06-28T16:22:00.637101Z","shell.execute_reply.started":"2021-06-28T16:21:59.976343Z","shell.execute_reply":"2021-06-28T16:22:00.633535Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"---\n## Train","metadata":{"papermill":{"duration":0.052353,"end_time":"2020-12-27T17:10:37.964527","exception":false,"start_time":"2020-12-27T17:10:37.912174","status":"completed"},"tags":[]}},{"cell_type":"code","source":"# LightningModule  --------------------------------------------------------------\nmodel = CycleGAN_LightningSystem(G_basestyle, G_stylebase, D_base, D_style, \n                                 lr, transform, reconstr_w, id_w)\n\n# Trainer  --------------------------------------------------------------\ntrainer = Trainer(\n    logger=False,\n    max_epochs=epoch,\n    gpus=1,\n    checkpoint_callback=False,\n    reload_dataloaders_every_epoch=True,\n    num_sanity_val_steps=0,  # Skip Sanity Check\n)\n\n\n# Train\nfor monet_limit in monet_limits:\n    dm = MonetDataModule(data_dir, transform, batch_size, monet_limit=monet_limit, seed=seed)\n    \n    trainer.fit(model, datamodule=dm)","metadata":{"papermill":{"duration":12852.714961,"end_time":"2020-12-27T20:44:50.732845","exception":false,"start_time":"2020-12-27T17:10:38.017884","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2021-06-28T16:09:18.031056Z","iopub.status.idle":"2021-06-28T16:09:18.031627Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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        img = img.cuda()\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')","metadata":{"papermill":{"duration":0.461003,"end_time":"2020-12-27T20:44:51.637573","exception":false,"start_time":"2020-12-27T20:44:51.17657","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2021-06-28T16:09:18.033156Z","iopub.status.idle":"2021-06-28T16:09:18.033718Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"submit(model, transform)","metadata":{"papermill":{"duration":1480.912732,"end_time":"2020-12-27T21:09:32.98677","exception":false,"start_time":"2020-12-27T20:44:52.074038","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2021-06-28T16:09:18.035038Z","iopub.status.idle":"2021-06-28T16:09:18.035604Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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()","metadata":{"papermill":{"duration":1.003672,"end_time":"2020-12-27T21:09:34.49409","exception":false,"start_time":"2020-12-27T21:09:33.490418","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2021-06-28T16:09:18.036892Z","iopub.status.idle":"2021-06-28T16:09:18.037455Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{"papermill":{"duration":0.457191,"end_time":"2020-12-27T21:09:35.59353","exception":false,"start_time":"2020-12-27T21:09:35.136339","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]}]}