{"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 torchsummary","metadata":{"execution":{"iopub.status.busy":"2022-12-17T13:40:19.770798Z","iopub.execute_input":"2022-12-17T13:40:19.771222Z","iopub.status.idle":"2022-12-17T13:40:27.940865Z","shell.execute_reply.started":"2022-12-17T13:40:19.771180Z","shell.execute_reply":"2022-12-17T13:40:27.939966Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip install pytorch_ssim","metadata":{"execution":{"iopub.status.busy":"2022-12-17T13:40:27.943088Z","iopub.execute_input":"2022-12-17T13:40:27.943473Z","iopub.status.idle":"2022-12-17T13:40:36.201576Z","shell.execute_reply.started":"2022-12-17T13:40:27.943432Z","shell.execute_reply":"2022-12-17T13:40:36.200629Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%matplotlib inline\nimport os, glob, random\nimport pytorch_ssim\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nfrom PIL import Image\nimport shutil\nfrom ignite.metrics import SSIM\nimport torch\nfrom random import sample\nfrom torchsummary import summary\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 sklearn.model_selection import train_test_split\nfrom pytorch_lightning import Trainer\nfrom pytorch_lightning.callbacks import EarlyStopping","metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","execution":{"iopub.status.busy":"2022-12-17T13:40:36.204440Z","iopub.execute_input":"2022-12-17T13:40:36.204799Z","iopub.status.idle":"2022-12-17T13:40:39.694306Z","shell.execute_reply.started":"2022-12-17T13:40:36.204753Z","shell.execute_reply":"2022-12-17T13:40:39.693135Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!wandb login 3e8c2e01bacb844caa27f15d0c4ae7f63d683a5e","metadata":{"execution":{"iopub.status.busy":"2022-12-17T13:40:39.696185Z","iopub.execute_input":"2022-12-17T13:40:39.696638Z","iopub.status.idle":"2022-12-17T13:40:42.118088Z","shell.execute_reply.started":"2022-12-17T13:40:39.696591Z","shell.execute_reply":"2022-12-17T13:40:42.117228Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"pl.__version__","metadata":{"execution":{"iopub.status.busy":"2022-12-17T13:40:42.122831Z","iopub.execute_input":"2022-12-17T13:40:42.123120Z","iopub.status.idle":"2022-12-17T13:40:42.129792Z","shell.execute_reply.started":"2022-12-17T13:40:42.123090Z","shell.execute_reply":"2022-12-17T13:40:42.128870Z"},"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-12-17T13:40:42.132306Z","iopub.execute_input":"2022-12-17T13:40:42.132915Z","iopub.status.idle":"2022-12-17T13:40:42.143391Z","shell.execute_reply.started":"2022-12-17T13:40:42.132876Z","shell.execute_reply":"2022-12-17T13:40:42.142608Z"},"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# Custom Dataset ---------------------------------------------------------------------------\nclass CustomDataset(Dataset):\n    def __init__(self, sketch_img_paths, photo_img_paths,  transform, phase='train'):\n        self.sketch_img_paths = sketch_img_paths\n        self.photo_img_paths = photo_img_paths\n        self.transform = transform\n        self.phase = phase\n\n    def __len__(self):\n        return min([len(self.sketch_img_paths), len(self.photo_img_paths)])\n\n    def __getitem__(self, idx):        \n        sketch_img_paths = self.sketch_img_paths[idx]\n        photo_idx = np.random.randint(0, len(self.photo_img_paths))\n        photo_img_paths = self.photo_img_paths[photo_idx]\n        sketch_img = Image.open(sketch_img_paths)\n        photo_img = Image.open(photo_img_paths)\n\n        sketch_img = self.transform(sketch_img, self.phase)\n        photo_img = self.transform(photo_img, self.phase)\n\n        return sketch_img, photo_img","metadata":{"execution":{"iopub.status.busy":"2022-12-17T13:40:42.144846Z","iopub.execute_input":"2022-12-17T13:40:42.145378Z","iopub.status.idle":"2022-12-17T13:40:42.159770Z","shell.execute_reply.started":"2022-12-17T13:40:42.145340Z","shell.execute_reply":"2022-12-17T13:40:42.159064Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Data Module\nclass DataModule(pl.LightningDataModule):\n    def __init__(self, data_dir, transform, batch_size, phase='train', seed=0):\n        super(DataModule, 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.sketch_img_paths = glob.glob(os.path.join(self.data_dir, 'sketches/sketches', '*.jpg'))\n        self.photo_img_paths = glob.glob(os.path.join(self.data_dir, 'photos/photos', '*.jpg'))\n        X_train, X_test, y_train, y_test = train_test_split(self.sketch_img_paths, self.photo_img_paths, test_size=0.2, random_state=42, shuffle=True)        \n        self.train_dataset = CustomDataset(X_train, y_train, self.transform, self.phase)\n        self.val_dataset = CustomDataset(X_test, y_test, self.transform, self.phase)\n        self.test_dataset = CustomDataset(X_test, y_test, self.transform, self.phase)\n\n    def train_dataloader(self):\n        '''\n        DataLoader của hàm train\n        '''\n        return DataLoader(self.train_dataset,\n                  batch_size=self.batch_size,\n                  shuffle=True,\n                  pin_memory=True,\n                  drop_last=True\n                 )\n    \n\n    def val_dataloader(self):\n        '''\n        DataLoader của hàm validate\n        '''\n        return DataLoader(\n            self.val_dataset,\n            batch_size=self.batch_size,\n        )\n\n    def test_dataloader(self):\n        '''\n        DataLoader của hàm test\n        '''\n        return DataLoader(\n            self.test_dataset,\n            batch_size=self.batch_size,\n        )","metadata":{"execution":{"iopub.status.busy":"2022-12-17T13:40:42.160949Z","iopub.execute_input":"2022-12-17T13:40:42.161218Z","iopub.status.idle":"2022-12-17T13:40:42.174436Z","shell.execute_reply.started":"2022-12-17T13:40:42.161184Z","shell.execute_reply":"2022-12-17T13:40:42.173519Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Sanity Check\ndata_dir = '/kaggle/input/cuhk-dataset-256'\ntransform = ImageTransform(img_size=256)\nbatch_size = 8\n\ndm = DataModule(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-12-17T13:40:42.175713Z","iopub.execute_input":"2022-12-17T13:40:42.176058Z","iopub.status.idle":"2022-12-17T13:40:46.881634Z","shell.execute_reply.started":"2022-12-17T13:40:42.176020Z","shell.execute_reply":"2022-12-17T13:40:46.880637Z"},"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-12-17T13:40:46.883350Z","iopub.execute_input":"2022-12-17T13:40:46.884072Z","iopub.status.idle":"2022-12-17T13:40:47.286244Z","shell.execute_reply.started":"2022-12-17T13:40:46.884028Z","shell.execute_reply":"2022-12-17T13:40:47.285320Z"},"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":{"execution":{"iopub.status.busy":"2022-12-17T13:40:47.287434Z","iopub.execute_input":"2022-12-17T13:40:47.287876Z","iopub.status.idle":"2022-12-17T13:40:47.645617Z","shell.execute_reply.started":"2022-12-17T13:40:47.287837Z","shell.execute_reply":"2022-12-17T13:40:47.644567Z"},"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-12-17T13:40:47.647059Z","iopub.execute_input":"2022-12-17T13:40:47.647400Z","iopub.status.idle":"2022-12-17T13:40:47.677996Z","shell.execute_reply.started":"2022-12-17T13:40:47.647367Z","shell.execute_reply":"2022-12-17T13:40:47.677260Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Sanity Check\n# net = CycleGAN_Unet_Generator()\n\n# out = net(base)\n# print(out.size())\nsummary(CycleGAN_Unet_Generator().float().cuda(), (3, 256, 256))","metadata":{"execution":{"iopub.status.busy":"2022-12-17T13:40:47.679204Z","iopub.execute_input":"2022-12-17T13:40:47.679715Z","iopub.status.idle":"2022-12-17T13:40:48.700828Z","shell.execute_reply.started":"2022-12-17T13:40:47.679662Z","shell.execute_reply":"2022-12-17T13:40:48.700025Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Sanity Check\n# net = CycleGAN_Discriminator()\n\n# out = net(base)\n# print(out.size())\nsummary(CycleGAN_Discriminator().float().cuda(), (3, 256, 256))","metadata":{"execution":{"iopub.status.busy":"2022-12-17T13:40:48.703933Z","iopub.execute_input":"2022-12-17T13:40:48.704218Z","iopub.status.idle":"2022-12-17T13:40:48.746737Z","shell.execute_reply.started":"2022-12-17T13:40:48.704181Z","shell.execute_reply":"2022-12-17T13:40:48.745926Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# CycleGAN - Lightning Module ---------------------------------------------------------------------------\nfrom ignite.engine.engine import Engine\nimport tensorflow as tf\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        self.metric_SSIM = []\n        \n        self.default_evaluator = Engine(self.eval_step)\n        self.metric = SSIM(data_range=1.0)\n        self.metric.attach(self.default_evaluator, 'ssim')\n        self.dm = DataModule(data_dir, transform, batch_size, seed=42) \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 eval_step(self, engine, batch):\n        return batch\n\n    def training_step(self, batch, batch_idx, optimizer_idx):\n        sketch_img, photo_img = batch\n        b = sketch_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(photo_img)), valid)\n            val_style = self.generator_loss(self.D_style(self.G_basestyle(sketch_img)), valid)\n            val_loss = (val_base + val_style) / 2\n\n            # Reconstruction\n            reconstr_base = self.mae(self.G_stylebase(self.G_basestyle(sketch_img)), sketch_img)\n            reconstr_style = self.mae(self.G_basestyle(self.G_stylebase(photo_img)), photo_img)\n            reconstr_loss = (reconstr_base + reconstr_style) / 2\n            \n            # Identity\n            id_base = self.mae(self.G_stylebase(sketch_img), sketch_img)\n            id_style = self.mae(self.G_basestyle(photo_img), photo_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            # SSIM metric\n            y_pre = self.G_basestyle(sketch_img)\n            state = self.default_evaluator.run([[y_pre, photo_img]]).metrics['ssim']\n            self.log('G_loss', G_loss)\n            self.log('SSIM_train', state)\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(photo_img)), fake)\n            D_style_gen_loss = self.discriminator_loss(self.D_style(self.G_basestyle(sketch_img)), fake)\n            D_base_valid_loss = self.discriminator_loss(self.D_base(sketch_img), valid)\n            D_style_valid_loss = self.discriminator_loss(self.D_style(photo_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            self.log('D_loss', D_loss)\n            return {'loss': D_loss}\n        \n    def validation_step(self, batch, batch_idx):\n            '''\n            Hàm này nhận vào một batch\n            lấy kết quả của mô hình và tính loss\n\n            '''\n            sketch_img, photo_img = batch\n            outputs = self.G_basestyle(sketch_img)\n            ssim = self.default_evaluator.run([[outputs, photo_img]]).metrics['ssim']\n            self.log(\"SSIM_train\", ssim)\n            self.log(\"SSIM_valid\", ssim)\n            return ssim\n\n    def test_step(self, batch, batch_idx):\n            sketch_img, photo_img = batch\n            outputs = self.G_basestyle(sketch_img)\n            ssim = self.default_evaluator.run([[outputs, photo_img]]).metrics['ssim']\n            self.log(\"SSIM_train\", ssim)\n            self.log(\"SSIM_test\", ssim)\n            return ssim\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#         ssim = sum([torch.stack([x['metric_SSIM'] 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.metric_SSIM.append(ssim)\n        \n        if self.step % 1 == 0:\n            # Display Model Output\n            target_img_paths = glob.glob('/kaggle/input/cuhk-dataset-256/sketches/sketches/*.jpg')\n            sample_data = sample(target_img_paths, 4)\n            target_imgs = [self.transform(Image.open(path), phase='test') for path in sample_data]\n            photos_imgs = [self.transform(Image.open(path.replace('sketches', 'photos')), phase='test') for path in sample_data]\n            target_imgs = torch.stack(target_imgs, dim=0)\n            target_imgs = target_imgs.cuda()\n            \n            photos_imgs = torch.stack(photos_imgs, dim=0)\n            photos_imgs = photos_imgs.cuda()\n            \n            gen_imgs = self.G_basestyle(target_imgs)\n            gen_img = torch.cat([target_imgs, gen_imgs, photos_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-12-17T13:40:48.748858Z","iopub.execute_input":"2022-12-17T13:40:48.749221Z","iopub.status.idle":"2022-12-17T13:40:53.280868Z","shell.execute_reply.started":"2022-12-17T13:40:48.749183Z","shell.execute_reply":"2022-12-17T13:40:53.279282Z"},"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='fa n_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-12-17T13:40:53.282844Z","iopub.execute_input":"2022-12-17T13:40:53.283209Z","iopub.status.idle":"2022-12-17T13:40:53.508390Z","shell.execute_reply.started":"2022-12-17T13:40:53.283133Z","shell.execute_reply":"2022-12-17T13:40:53.507322Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"---\n## Train","metadata":{}},{"cell_type":"code","source":"wandb_logger = pl.loggers.WandbLogger(\n    project=\"GAN-project\", name='model-gan-ver-3', log_model=\"all\"\n)\nlr_monitor = pl.callbacks.LearningRateMonitor(logging_interval=\"step\")","metadata":{"execution":{"iopub.status.busy":"2022-12-17T13:40:53.513098Z","iopub.execute_input":"2022-12-17T13:40:53.515577Z","iopub.status.idle":"2022-12-17T13:40:53.527801Z","shell.execute_reply.started":"2022-12-17T13:40:53.515536Z","shell.execute_reply":"2022-12-17T13:40:53.526904Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Config  -----------------------------------------------------------------\ndata_dir = '/kaggle/input/cuhk-dataset-256'\ntransform = ImageTransform(img_size=256)\nbatch_size = 1\nlr = {\n    'G': 0.0002,\n    'D': 0.0002\n}\nepoch = 50\nseed = 42\nreconstr_w = 10\nid_w = 2\nseed_everything(seed)\n\n# DataModule  -----------------------------------------------------------------\ndm = DataModule(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\nlr_monitor = pl.callbacks.LearningRateMonitor(logging_interval=\"step\")\n# Trainer  --------------------------------------------------------------\ntrainer = Trainer(\n    logger=wandb_logger,\n    callbacks=[lr_monitor],\n    max_epochs=epoch,\n    gpus=1,\n    checkpoint_callback=False,\n    log_every_n_steps=1,\n    reload_dataloaders_every_epoch=True\n)\n\n\n# Train\ntrainer.fit(model, datamodule=dm)","metadata":{"execution":{"iopub.status.busy":"2022-12-17T13:40:53.532944Z","iopub.execute_input":"2022-12-17T13:40:53.535683Z","iopub.status.idle":"2022-12-17T13:57:48.506961Z","shell.execute_reply.started":"2022-12-17T13:40:53.535644Z","shell.execute_reply":"2022-12-17T13:57:48.505979Z"},"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('/kaggle/input/cuhk-dataset-256/sketches/sketches/*.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#         5      \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":{"execution":{"iopub.status.busy":"2022-12-17T13:57:48.508856Z","iopub.execute_input":"2022-12-17T13:57:48.509228Z","iopub.status.idle":"2022-12-17T13:57:48.514689Z","shell.execute_reply.started":"2022-12-17T13:57:48.509190Z","shell.execute_reply":"2022-12-17T13:57:48.513872Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# submit(model, transform)","metadata":{"execution":{"iopub.status.busy":"2022-12-17T13:57:48.515904Z","iopub.execute_input":"2022-12-17T13:57:48.516412Z","iopub.status.idle":"2022-12-17T13:57:48.529000Z","shell.execute_reply.started":"2022-12-17T13:57:48.516376Z","shell.execute_reply":"2022-12-17T13:57:48.528316Z"},"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-12-17T13:57:48.532002Z","iopub.execute_input":"2022-12-17T13:57:48.532275Z","iopub.status.idle":"2022-12-17T13:57:48.852703Z","shell.execute_reply.started":"2022-12-17T13:57:48.532251Z","shell.execute_reply":"2022-12-17T13:57:48.851854Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"torch.save(model.state_dict(), \"/kaggle/working/model_ver4.pth\")","metadata":{"execution":{"iopub.status.busy":"2022-12-17T13:57:48.854118Z","iopub.execute_input":"2022-12-17T13:57:48.854707Z","iopub.status.idle":"2022-12-17T13:57:49.392419Z","shell.execute_reply.started":"2022-12-17T13:57:48.854667Z","shell.execute_reply":"2022-12-17T13:57:49.391592Z"},"trusted":true},"execution_count":null,"outputs":[]}]}