training-neongan.ipynb

{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.7.6","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"}},"nbformat_minor":4,"nbformat":4,"cells":[{"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\nfrom pytorch_lightning.callbacks import ModelCheckpoint","metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","execution":{"iopub.status.busy":"2022-04-25T09:49:50.145593Z","iopub.execute_input":"2022-04-25T09:49:50.145938Z","iopub.status.idle":"2022-04-25T09:49:52.574493Z","shell.execute_reply.started":"2022-04-25T09:49:50.145907Z","shell.execute_reply":"2022-04-25T09:49:52.573466Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"pl.__version__","metadata":{"execution":{"iopub.status.busy":"2022-04-25T09:44:02.121326Z","iopub.execute_input":"2022-04-25T09:44:02.121639Z","iopub.status.idle":"2022-04-25T09:44:02.131315Z","shell.execute_reply.started":"2022-04-25T09:44:02.121602Z","shell.execute_reply":"2022-04-25T09:44:02.13041Z"},"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":{"execution":{"iopub.status.busy":"2022-04-25T09:49:58.963825Z","iopub.execute_input":"2022-04-25T09:49:58.964207Z","iopub.status.idle":"2022-04-25T09:49:58.970819Z","shell.execute_reply.started":"2022-04-25T09:49:58.964168Z","shell.execute_reply":"2022-04-25T09:49:58.970042Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"---\n## Dataset","metadata":{}},{"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.RandomHorizontalFlip(),\n                transforms.RandomVerticalFlip(),\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# neon Dataset ---------------------------------------------------------------------------\nclass neonDataset(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","metadata":{"execution":{"iopub.status.busy":"2022-04-25T09:50:07.321093Z","iopub.execute_input":"2022-04-25T09:50:07.321667Z","iopub.status.idle":"2022-04-25T09:50:07.335925Z","shell.execute_reply.started":"2022-04-25T09:50:07.321631Z","shell.execute_reply":"2022-04-25T09:50:07.334961Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Data Module\nclass neonDataModule(pl.LightningDataModule):\n    def __init__(self, data_dir, transform, batch_size, phase='train', seed=0):\n        super(neonDataModule, 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\n    def prepare_data(self):\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('../input/futuristic-images/Image_Dataset_256/', '*.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 = neonDataset(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":{"execution":{"iopub.status.busy":"2022-04-25T09:50:08.483236Z","iopub.execute_input":"2022-04-25T09:50:08.483585Z","iopub.status.idle":"2022-04-25T09:50:08.494497Z","shell.execute_reply.started":"2022-04-25T09:50:08.483554Z","shell.execute_reply":"2022-04-25T09:50:08.493521Z"},"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 = neonDataModule(data_dir, transform, batch_size, phase='test')\ndm.prepare_data()\n\ndataloader = dm.train_dataloader()\nbase, style = next(iter(dataloader))\n\nprint('Input Shape {}, {}'.format(base.size(), style.size()))","metadata":{"execution":{"iopub.status.busy":"2022-04-25T09:50:09.764457Z","iopub.execute_input":"2022-04-25T09:50:09.765014Z","iopub.status.idle":"2022-04-25T09:50:10.389886Z","shell.execute_reply.started":"2022-04-25T09:50:09.764966Z","shell.execute_reply":"2022-04-25T09:50:10.389006Z"},"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":{"execution":{"iopub.status.busy":"2022-04-25T09:50:10.614578Z","iopub.execute_input":"2022-04-25T09:50:10.614975Z","iopub.status.idle":"2022-04-25T09:50:11.085981Z","shell.execute_reply.started":"2022-04-25T09:50:10.614937Z","shell.execute_reply":"2022-04-25T09:50:11.085262Z"},"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('Future Pictures')\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2022-04-25T09:50:12.260928Z","iopub.execute_input":"2022-04-25T09:50:12.261667Z","iopub.status.idle":"2022-04-25T09:50:12.683687Z","shell.execute_reply.started":"2022-04-25T09:50:12.261617Z","shell.execute_reply":"2022-04-25T09:50:12.682601Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"---\n## Model","metadata":{}},{"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            Downsample(filter * 8, filter * 8), # (b, filter * 8, 4, 4)\n            Downsample(filter * 8, filter * 8), # (b, filter * 8, 2, 2)\n        ])\n\n        self.upsamples = nn.ModuleList([\n            Upsample(filter * 8, filter * 8),\n            Upsample(filter * 16, filter * 8),\n            Upsample(filter * 16, 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":{"execution":{"iopub.status.busy":"2022-04-25T09:50:15.376836Z","iopub.execute_input":"2022-04-25T09:50:15.377163Z","iopub.status.idle":"2022-04-25T09:50:15.411154Z","shell.execute_reply.started":"2022-04-25T09:50:15.377135Z","shell.execute_reply":"2022-04-25T09:50:15.40976Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Sanity Check\nnet = CycleGAN_Unet_Generator()\n\nout = net(base)\nprint(out.size())","metadata":{"execution":{"iopub.status.busy":"2022-04-25T09:50:16.59248Z","iopub.execute_input":"2022-04-25T09:50:16.593097Z","iopub.status.idle":"2022-04-25T09:50:18.69443Z","shell.execute_reply.started":"2022-04-25T09:50:16.593058Z","shell.execute_reply":"2022-04-25T09:50:18.693611Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Sanity Check\nnet = CycleGAN_Discriminator()\n\nout = net(base)\nprint(out.size())","metadata":{"execution":{"iopub.status.busy":"2022-04-25T09:50:18.697401Z","iopub.execute_input":"2022-04-25T09:50:18.697908Z","iopub.status.idle":"2022-04-25T09:50:19.339763Z","shell.execute_reply.started":"2022-04-25T09:50:18.697875Z","shell.execute_reply":"2022-04-25T09:50:19.338797Z"},"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, log_dict = {}):\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        self.log_dict = {}\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            values = {'loss': G_loss, 'validity': val_loss, 'reconstr': reconstr_loss, 'identity': id_loss}\n            #self.log_dict(values, on_step=False, on_epoch=True,logger=True)\n            \n            self.log('validity',val_loss)\n            \n            return {'loss': G_loss, 'validity': val_loss, 'reconstr': reconstr_loss, 'identity': id_loss}\n            \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 self.step % 30 == 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":{"execution":{"iopub.status.busy":"2022-04-25T09:50:19.341059Z","iopub.execute_input":"2022-04-25T09:50:19.341319Z","iopub.status.idle":"2022-04-25T09:50:19.387277Z","shell.execute_reply.started":"2022-04-25T09:50:19.341294Z","shell.execute_reply":"2022-04-25T09:50:19.386354Z"},"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":{"execution":{"iopub.status.busy":"2022-04-25T09:50:19.404354Z","iopub.execute_input":"2022-04-25T09:50:19.404862Z","iopub.status.idle":"2022-04-25T09:50:19.418954Z","shell.execute_reply.started":"2022-04-25T09:50:19.404826Z","shell.execute_reply":"2022-04-25T09:50:19.418296Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"---\n## Train","metadata":{}},{"cell_type":"code","source":"# Config  -----------------------------------------------------------------\ndata_dir = '../input/gan-getting-started'\ntransform = ImageTransform(img_size=256)\nbatch_size = 1\nlr = {\n    'G': 0.0002,\n    'D': 0.0002\n}\nepoch = 180\nseed = 42\nreconstr_w = 10\nid_w = 2\nseed_everything(seed)\n\n# DataModule  -----------------------------------------------------------------\ndm = neonDataModule(data_dir, transform, batch_size, seed=seed)\n\nG_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')\n\n# LightningModule  --------------------------------------------------------------\nmodel = CycleGAN_LightningSystem(G_basestyle, G_stylebase, D_base, D_style, \n                                 lr, transform, reconstr_w, id_w)\n\n#checkpoint_callback = ModelCheckpoint(filepath='./-{epoch:02d}-{val_cer:.2f}',save_last=True, mode='min', monitor='loss', save_top_k=5)\n\n   \n\n# Trainer  --------------------------------------------------------------\ntrainer = Trainer(\n    logger=True,\n    max_epochs=epoch,\n    gpus=1,\n    reload_dataloaders_every_epoch=True,\n    num_sanity_val_steps=0,  # Skip Sanity Check\n    default_root_dir=\"./\",\n    \n)\n\n\n\n# Train\ntrainer.fit(model, datamodule=dm)\n#print(checkpoint_callback.best_model_path, checkpoint_callback.best_model_score)","metadata":{"execution":{"iopub.status.busy":"2022-04-24T19:53:27.042588Z","iopub.execute_input":"2022-04-24T19:53:27.04304Z","iopub.status.idle":"2022-04-24T19:57:24.439798Z","shell.execute_reply.started":"2022-04-24T19:53:27.042979Z","shell.execute_reply":"2022-04-24T19:57:24.43896Z"},"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":{"execution":{"iopub.status.busy":"2022-04-24T19:58:16.276358Z","iopub.execute_input":"2022-04-24T19:58:16.276725Z","iopub.status.idle":"2022-04-24T19:58:16.598628Z","shell.execute_reply.started":"2022-04-24T19:58:16.276693Z","shell.execute_reply":"2022-04-24T19:58:16.597696Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"---\n## Load Checkpoint","metadata":{}},{"cell_type":"code","source":"transform = ImageTransform(img_size=256)\nbatch_size = 1\nlr = {\n    'G': 0.0002,\n    'D': 0.0002\n}\nepoch = 180\nseed = 42\nreconstr_w = 10\nid_w = 2\nseed_everything(seed)\n\n# DataModule  -----------------------------------------------------------------\ndm = neonDataModule(data_dir, transform, batch_size, seed=seed)\n\nG_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')\n\n# LightningModule  --------------------------------------------------------------\nmodel = CycleGAN_LightningSystem(G_basestyle, G_stylebase, D_base, D_style, \n                                 lr, transform, reconstr_w, id_w)","metadata":{"execution":{"iopub.status.busy":"2022-04-25T09:50:33.071062Z","iopub.execute_input":"2022-04-25T09:50:33.071412Z","iopub.status.idle":"2022-04-25T09:50:34.606246Z","shell.execute_reply.started":"2022-04-25T09:50:33.071378Z","shell.execute_reply":"2022-04-25T09:50:34.605507Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"checkpoint = torch.load('../input/neongan-model/epoch179.ckpt')","metadata":{"execution":{"iopub.status.busy":"2022-04-25T08:05:06.357663Z","iopub.execute_input":"2022-04-25T08:05:06.357972Z","iopub.status.idle":"2022-04-25T08:05:14.537227Z","shell.execute_reply.started":"2022-04-25T08:05:06.357944Z","shell.execute_reply":"2022-04-25T08:05:14.536448Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model.load_state_dict(checkpoint['state_dict'])","metadata":{"execution":{"iopub.status.busy":"2022-04-25T08:05:30.273745Z","iopub.execute_input":"2022-04-25T08:05:30.274052Z","iopub.status.idle":"2022-04-25T08:05:30.359142Z","shell.execute_reply.started":"2022-04-25T08:05:30.274023Z","shell.execute_reply":"2022-04-25T08:05:30.358476Z"},"trusted":true},"execution_count":null,"outputs":[]}]}