{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.10","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":21755,"databundleVersionId":1475600,"sourceType":"competition"},{"sourceId":3956508,"sourceType":"datasetVersion","datasetId":850761}],"dockerImageVersionId":30498,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# **Photos To Monet**","metadata":{}},{"cell_type":"code","source":"!pip install deepspeed","metadata":{"execution":{"iopub.status.busy":"2023-12-14T04:55:36.108097Z","iopub.execute_input":"2023-12-14T04:55:36.108339Z","iopub.status.idle":"2023-12-14T04:56:03.035186Z","shell.execute_reply.started":"2023-12-14T04:55:36.108315Z","shell.execute_reply":"2023-12-14T04:56:03.034017Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import glob\nimport os\nimport shutil\nimport deepspeed as ds\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport pytorch_lightning as L\nimport torch\nimport torch.nn.functional as F\nimport torchvision.transforms as T\nfrom pytorch_lightning.utilities import CombinedLoader\nfrom torch import nn\nfrom torch.utils.data import DataLoader, Dataset\nfrom torchvision.io import read_image\nfrom torchvision.utils import make_grid, save_image\n_ = L.seed_everything(0, workers=True)","metadata":{"execution":{"iopub.status.busy":"2023-12-14T04:56:03.037608Z","iopub.execute_input":"2023-12-14T04:56:03.037991Z","iopub.status.idle":"2023-12-14T04:56:22.858570Z","shell.execute_reply.started":"2023-12-14T04:56:03.037954Z","shell.execute_reply":"2023-12-14T04:56:22.857673Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Data Preprocessing ","metadata":{}},{"cell_type":"code","source":"def show_img(img_tensor, nrow, title=\"\"):\n    img_tensor = img_tensor.detach().cpu() * 0.5 + 0.5\n    img_grid = make_grid(img_tensor, nrow=nrow).permute(1, 2, 0)\n    plt.figure(figsize=(10, 7))\n    plt.imshow(img_grid)\n    plt.axis(\"off\")\n    plt.title(title)\n    plt.show()","metadata":{"execution":{"iopub.status.busy":"2023-12-14T04:56:22.859739Z","iopub.execute_input":"2023-12-14T04:56:22.860057Z","iopub.status.idle":"2023-12-14T04:56:22.866005Z","shell.execute_reply.started":"2023-12-14T04:56:22.860031Z","shell.execute_reply":"2023-12-14T04:56:22.865036Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Augmenting the images.**\n\nBefore loading the datasets, we define CustomTransform for image augmentation. This improves learning by introducing more variety in the images during training instead of learning from the same set of images, especially when we only have 300 Monet paintings.","metadata":{}},{"cell_type":"code","source":"class CustomTransform(object):\n    def __init__(self, load_dim=286, target_dim=256):\n        self.transform_train = T.Compose([\n            T.Resize((load_dim, load_dim), antialias=True),\n            T.RandomCrop((target_dim, target_dim)),\n            T.RandomHorizontalFlip(p=0.5),\n            T.ColorJitter(brightness=0.2, contrast=0.2,\n                          saturation=0.2, hue=0.1),\n        ])\n        self.transform = T.Resize((target_dim, target_dim), antialias=True)   \n    def __call__(self, img, stage):\n        if stage == \"fit\":\n            img = self.transform_train(img)\n        else:\n            img = self.transform(img)\n        return img * 2 - 1","metadata":{"execution":{"iopub.status.busy":"2023-12-14T04:56:22.868244Z","iopub.execute_input":"2023-12-14T04:56:22.868570Z","iopub.status.idle":"2023-12-14T04:56:22.885194Z","shell.execute_reply.started":"2023-12-14T04:56:22.868544Z","shell.execute_reply":"2023-12-14T04:56:22.884183Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Storing dataset**","metadata":{}},{"cell_type":"code","source":"class CustomDataset(Dataset):\n    def __init__(self, filenames, transform, stage):\n        self.filenames = filenames\n        self.transform = transform\n        self.stage = stage\n        \n    def __len__(self):\n        return len(self.filenames)\n    \n    def __getitem__(self, idx):\n        img_name = self.filenames[idx]\n        img = read_image(img_name) / 255.0\n        return self.transform(img, stage=self.stage)","metadata":{"execution":{"iopub.status.busy":"2023-12-14T04:56:22.886717Z","iopub.execute_input":"2023-12-14T04:56:22.887041Z","iopub.status.idle":"2023-12-14T04:56:22.896035Z","shell.execute_reply.started":"2023-12-14T04:56:22.887016Z","shell.execute_reply":"2023-12-14T04:56:22.895165Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Iterating through the datasets**\n\nTo prepare the datasets, we load them into DataLoader separately, which can then iterate through the datasets as needed. Because the training dataset contains both the Monet paintings and photos, we pass both dataloaders into CombinedLoader for training. ","metadata":{}},{"cell_type":"code","source":"DEBUG = False\n\nDM_CONFIG = {    \n#     \"monet_dir\": os.path.join(\"/kaggle/input/van-gogh-paintings\", \"**\", \"*.jpg\"),\n    \"monet_dir\": \"/kaggle/input/van-gogh-paintings\",\n    \"photo_dir\": os.path.join(\"/kaggle/input/gan-getting-started/photo_jpg\", \"*.jpg\"),\n    \n    \"loader_config\": {\n        \"num_workers\": os.cpu_count(),\n        \"pin_memory\": torch.cuda.is_available(),\n    },\n    \"sample_size\": 5,\n    \"batch_size\": 1 if not DEBUG else 1,\n}","metadata":{"execution":{"iopub.status.busy":"2023-12-14T04:56:22.897029Z","iopub.execute_input":"2023-12-14T04:56:22.897304Z","iopub.status.idle":"2023-12-14T04:56:22.917235Z","shell.execute_reply.started":"2023-12-14T04:56:22.897283Z","shell.execute_reply":"2023-12-14T04:56:22.916431Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_limited_image_paths(directory, limit_per_folder=40, specific_folders=None):\n    all_paths = []\n    for root, dirs, files in os.walk(directory):\n        # Check if the current folder is one of the specific folders\n        if specific_folders is not None:\n            # Only proceed if the folder name matches one of the specific folders\n            if not any(specific_folder in root for specific_folder in specific_folders):\n                continue\n\n        image_paths = [os.path.join(root, file) for file in files if file.endswith(\".jpg\")]\n        all_paths.extend(image_paths[:limit_per_folder])\n    return all_paths\n\n\nclass CustomDataModule(L.LightningDataModule):\n    def __init__(\n        self,\n        monet_dir,\n        photo_dir, \n        loader_config,\n        sample_size,\n        batch_size,\n    ):\n        super().__init__()\n        self.loader_config = loader_config\n        self.sample_size = sample_size\n        self.batch_size = batch_size       \n        # Example of calling the function\n        specific_folders = ['Aries', 'Auvers sur Oise', 'Saint Remy', 'Nuenen']  # Replace with actual folder names or patterns\n        image_paths = get_limited_image_paths(monet_dir, limit_per_folder=189, specific_folders=specific_folders)\n\n        # store file paths\n#         self.monet_filenames = sorted(glob.glob(monet_dir, recursive=True))\n        self.monet_filenames = sorted(image_paths)\n        self.photo_filenames = sorted(glob.glob(photo_dir, recursive=True))[:2000]\n        print(f\"Number of Monet images: {len(self.monet_filenames)}\")\n        print(f\"Number of photo images: {len(self.photo_filenames)}\")\n\n        \n        # define transformations for image augmentation\n        self.transform = CustomTransform()\n        \n    def setup(self, stage):\n        if stage == \"fit\":\n            self.train_monet = CustomDataset(self.monet_filenames, self.transform, stage)\n            self.train_photo = CustomDataset(self.photo_filenames, self.transform, stage)\n            \n        if stage in [\"fit\", \"test\", \"predict\"]:\n            # to be used for test and prediction datasets as well\n            self.valid_photo = CustomDataset(self.photo_filenames, self.transform, None)\n            \n    def train_dataloader(self):\n        loader_config = {\n            \"shuffle\": True,\n            \"drop_last\": True,\n            \"batch_size\": self.batch_size,\n            **self.loader_config,\n        }\n        loader_monet = DataLoader(self.train_monet, **loader_config)\n        loader_photo = DataLoader(self.train_photo, **loader_config)\n        loaders = {\"monet\": loader_monet, \"photo\": loader_photo}\n        return CombinedLoader(loaders, mode=\"max_size_cycle\")\n    \n    def val_dataloader(self):\n        return DataLoader(\n            self.valid_photo,\n            batch_size=self.sample_size,\n            **self.loader_config,\n        )\n    \n    def test_dataloader(self):\n        return self.val_dataloader()\n    \n    def predict_dataloader(self):\n        return DataLoader(\n            self.valid_photo,\n            batch_size=self.batch_size,\n            **self.loader_config,\n        )","metadata":{"execution":{"iopub.status.busy":"2023-12-14T04:56:22.920054Z","iopub.execute_input":"2023-12-14T04:56:22.920526Z","iopub.status.idle":"2023-12-14T04:56:22.934673Z","shell.execute_reply.started":"2023-12-14T04:56:22.920502Z","shell.execute_reply":"2023-12-14T04:56:22.933725Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"We check that the datamodule defined is working as intended by visualizing samples of the images below.","metadata":{}},{"cell_type":"code","source":"dm_sample = CustomDataModule(batch_size=5, **{k: v for k, v in DM_CONFIG.items() if k != \"batch_size\"})\ndm_sample.setup(\"fit\")\ntrain_loader = dm_sample.train_dataloader()\nimgs = next(iter(train_loader))\nshow_img(imgs[\"monet\"], nrow=5, title=\"Augmented Monet Paintings\")\nshow_img(imgs[\"photo\"], nrow=5, title=\"Augmented Photos\")","metadata":{"execution":{"iopub.status.busy":"2023-12-14T04:56:22.935836Z","iopub.execute_input":"2023-12-14T04:56:22.936103Z","iopub.status.idle":"2023-12-14T04:56:30.938553Z","shell.execute_reply.started":"2023-12-14T04:56:22.936081Z","shell.execute_reply":"2023-12-14T04:56:30.937188Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Building GAN Architecture","metadata":{}},{"cell_type":"markdown","source":"**U-Net Generator**","metadata":{}},{"cell_type":"markdown","source":"A common architecture for the CycleGAN generator is the U-Net. U-Net is a network which consists of a sequence of downsampling blocks followed by a sequence of upsampling blocks, giving it the U-shaped architecture. In the upsampling path, we concatenate the outputs of the upsampling blocks and the outputs of the downsampling blocks symmetrically. This can be seen as a kind of skip connection, facilitating information flow in deep networks and reducing the impact of vanishing gradients.","metadata":{}},{"cell_type":"markdown","source":"**ResNet generator**","metadata":{}},{"cell_type":"markdown","source":"Similar to the U-Net architecture, the ResNet generator consists of the downsampling path and upsampling path. The difference is that the ResNet generator does not have the long skip connections from the concatenations of outputs. Instead, the ResNet generator uses residual blocks between the two paths. These residual blocks have short skip connections where the original input is added to the output.","metadata":{}},{"cell_type":"markdown","source":"**Downsampling blocks**","metadata":{}},{"cell_type":"markdown","source":"The downsampling blocks use convolution layers to increase the number of feature maps while reducing the dimensions of the 2D image.","metadata":{}},{"cell_type":"code","source":"class Downsampling(nn.Module):\n    def __init__(self, in_channels, out_channels, kernel_size=4, stride=2, padding=1, norm=True, lrelu=True):\n        super().__init__()\n        layers = [nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, bias=not norm)]\n        if norm:\n            layers.append(nn.InstanceNorm2d(out_channels))\n        if lrelu:\n            layers.append(nn.LeakyReLU(0.2, inplace=True))\n        else:\n            layers.append(nn.ReLU(inplace=True))\n        self.block = nn.Sequential(*layers)\n\n    def forward(self, x):\n        return self.block(x)","metadata":{"execution":{"iopub.status.busy":"2023-12-14T04:56:30.940496Z","iopub.execute_input":"2023-12-14T04:56:30.940891Z","iopub.status.idle":"2023-12-14T04:56:30.952333Z","shell.execute_reply.started":"2023-12-14T04:56:30.940839Z","shell.execute_reply":"2023-12-14T04:56:30.951177Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Upsampling blocks**\n\nOn the other hand, the upsampling blocks contain transposed convolution layers, which combine the learned features to output an image with the original size.","metadata":{}},{"cell_type":"code","source":"class Upsampling(nn.Module):\n    def __init__(self, in_channels, out_channels, kernel_size=4, stride=2, padding=1, output_padding=0, dropout=False):\n        super().__init__()\n        layers = [nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride, padding, output_padding, bias=False),\n                  nn.InstanceNorm2d(out_channels),\n                  nn.ReLU(inplace=True)]\n        if dropout:\n            layers.append(nn.Dropout(0.5))\n        self.block = nn.Sequential(*layers)\n\n    def forward(self, x):\n        return self.block(x)","metadata":{"execution":{"iopub.status.busy":"2023-12-14T04:56:30.957067Z","iopub.execute_input":"2023-12-14T04:56:30.958112Z","iopub.status.idle":"2023-12-14T04:56:30.967686Z","shell.execute_reply.started":"2023-12-14T04:56:30.958078Z","shell.execute_reply":"2023-12-14T04:56:30.966662Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Residual blocks**\n\nAs described above, the residual blocks have convolution layers where the original input is added to the output.","metadata":{}},{"cell_type":"code","source":"class ResNetBlock(nn.Module):\n    def __init__(self, dim):\n        super(ResNetBlock, self).__init__()\n        self.conv_block = self.build_conv_block(dim)\n\n    def build_conv_block(self, dim):\n        conv_block = [\n            nn.ReflectionPad2d(1),\n            nn.Conv2d(dim, dim, kernel_size=3, padding=0, bias=False),\n            nn.InstanceNorm2d(dim),\n            nn.ReLU(True),\n\n            nn.ReflectionPad2d(1),\n            nn.Conv2d(dim, dim, kernel_size=3, padding=0, bias=False),\n            nn.InstanceNorm2d(dim)\n        ]\n        return nn.Sequential(*conv_block)\n\n    def forward(self, x):\n        return x + self.conv_block(x)","metadata":{"execution":{"iopub.status.busy":"2023-12-14T04:56:30.969596Z","iopub.execute_input":"2023-12-14T04:56:30.970026Z","iopub.status.idle":"2023-12-14T04:56:30.978717Z","shell.execute_reply.started":"2023-12-14T04:56:30.969992Z","shell.execute_reply":"2023-12-14T04:56:30.977869Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Building the generator**\n\nWith the building blocks defined, we can now build our CycleGAN generator. For reference, the output size of each block is commented below.","metadata":{}},{"cell_type":"code","source":"class UNetGenerator(nn.Module):\n    def __init__(self, hid_channels, in_channels, out_channels):\n        super().__init__()\n        self.down1 = Downsampling(in_channels, hid_channels)\n        self.down2 = Downsampling(hid_channels, hid_channels * 2)\n        self.down3 = Downsampling(hid_channels * 2, hid_channels * 4)\n        self.down4 = Downsampling(hid_channels * 4, hid_channels * 8)\n\n        self.bottleneck = Downsampling(hid_channels * 8, hid_channels * 8, norm=False)\n\n        self.up1 = Upsampling(hid_channels * 8, hid_channels * 8, dropout=True)\n        self.up2 = Upsampling(hid_channels * 16, hid_channels * 4, dropout=True)\n        self.up3 = Upsampling(hid_channels * 8, hid_channels * 2)\n        self.up4 = Upsampling(hid_channels * 4, hid_channels)\n\n        self.out = nn.Sequential(\n            nn.ConvTranspose2d(hid_channels * 2, out_channels, 4, 2, 1),\n            nn.Tanh()\n        )\n\n    def forward(self, x):\n        # Downsampling\n        d1 = self.down1(x)\n        d2 = self.down2(d1)\n        d3 = self.down3(d2)\n        d4 = self.down4(d3)\n\n        bottleneck = self.bottleneck(d4)\n\n        # Upsampling with skip connections\n        u1 = self.up1(bottleneck)\n        u1 = torch.cat([u1, d4], dim=1)\n\n        u2 = self.up2(u1)\n        u2 = torch.cat([u2, d3], dim=1)\n\n        u3 = self.up3(u2)\n        u3 = torch.cat([u3, d2], dim=1)\n\n        u4 = self.up4(u3)\n        u4 = torch.cat([u4, d1], dim=1)\n\n        return self.out(u4)\n\n\nclass ResNetGenerator(nn.Module):\n    def __init__(self, in_channels, out_channels, hid_channels=64, n_blocks=9, use_dropout=False):\n        assert(n_blocks >= 0)\n        super(ResNetGenerator, self).__init__()\n\n        model = [nn.ReflectionPad2d(3),\n                 nn.Conv2d(in_channels, hid_channels, kernel_size=7, padding=0, bias=False),\n                 nn.InstanceNorm2d(hid_channels),\n                 nn.ReLU(True)]\n\n        # Downsampling\n        n_downsampling = 2\n        for i in range(n_downsampling):\n            mult = 2**i\n            model += [nn.Conv2d(hid_channels * mult, hid_channels * mult * 2, kernel_size=3, stride=2, padding=1, bias=False),\n                      nn.InstanceNorm2d(hid_channels * mult * 2),\n                      nn.ReLU(True)]\n\n        # ResNet Blocks\n        mult = 2**n_downsampling\n        for i in range(n_blocks):\n            model += [ResNetBlock(hid_channels * mult)]\n\n        # Upsampling\n        for i in range(n_downsampling):\n            mult = 2**(n_downsampling - i)\n            model += [nn.ConvTranspose2d(hid_channels * mult, int(hid_channels * mult / 2),\n                                         kernel_size=3, stride=2,\n                                         padding=1, output_padding=1,\n                                         bias=False),\n                      nn.InstanceNorm2d(int(hid_channels * mult / 2)),\n                      nn.ReLU(True)]\n\n        model += [nn.ReflectionPad2d(3)]\n        model += [nn.Conv2d(hid_channels, out_channels, kernel_size=7, padding=0)]\n        model += [nn.Tanh()]\n\n        self.model = nn.Sequential(*model)\n\n    def forward(self, input):\n        return self.model(input)\n    \n    \ndef get_gen(gen_name, hid_channels, num_resblocks, in_channels=3, out_channels=3):\n    if gen_name == \"unet\":\n        return UNetGenerator(hid_channels, in_channels, out_channels)\n    elif gen_name == \"resnet\":\n        return ResNetGenerator(in_channels, out_channels, hid_channels, num_resblocks)\n    else:\n        raise NotImplementedError(f\"Generator name '{gen_name}' not recognized.\")","metadata":{"execution":{"iopub.status.busy":"2023-12-14T04:56:30.980021Z","iopub.execute_input":"2023-12-14T04:56:30.980328Z","iopub.status.idle":"2023-12-14T04:56:31.002644Z","shell.execute_reply.started":"2023-12-14T04:56:30.980304Z","shell.execute_reply":"2023-12-14T04:56:31.001786Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Patch Gan Generator**\n\nUnlike conventional networks that output a single probability of the input image being real or fake, CycleGAN uses the PatchGAN discriminator that outputs a matrix of values. Intuitively, each value of the output matrix checks the corresponding portion of the input image. Values closer to 1 indicate real classification and values closer to 0 indicate fake classification.","metadata":{}},{"cell_type":"markdown","source":"**Building the discriminator**\n\nIn general, the PatchGAN discriminator consists of a sequence of convolution layers, which can be built using the downsampling blocks defined earlier.","metadata":{}},{"cell_type":"code","source":"class Discriminator(nn.Module):\n    def __init__(self, hid_channels, in_channels=3):\n        super().__init__()\n        self.block = nn.Sequential(\n            Downsampling(in_channels, hid_channels, norm=False), # 64x128x128\n            Downsampling(hid_channels, hid_channels*2), # 128x64x64\n            Downsampling(hid_channels*2, hid_channels*4), # 256x32x32\n            Downsampling(hid_channels*4, hid_channels*8, stride=1), # 512x31x31\n            nn.Conv2d(hid_channels*8, 1, kernel_size=4, padding=1), # 1x30x30\n        )\n        \n    def forward(self, x):\n        return self.block(x)","metadata":{"execution":{"iopub.status.busy":"2023-12-14T04:56:31.003866Z","iopub.execute_input":"2023-12-14T04:56:31.004111Z","iopub.status.idle":"2023-12-14T04:56:31.016376Z","shell.execute_reply.started":"2023-12-14T04:56:31.004090Z","shell.execute_reply":"2023-12-14T04:56:31.015666Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Image buffer**\n\nThe original CycleGAN implementation updates the discriminator using a history of generated images instead of the latest images generated. This is done by setting up an image buffer that stores previously generated images. With probability 0.5, each newly generated image is swapped with a previously generated image stored in the buffer. This stabilizes training by giving the discriminator access to past information.","metadata":{}},{"cell_type":"code","source":"class ImageBuffer(object):\n    def __init__(self, buffer_size):\n        self.buffer_size = buffer_size\n        if self.buffer_size > 0:\n            # the current capacity of the buffer\n            self.curr_cap = 0\n            # initialize buffer as empty list\n            self.buffer = []\n    \n    def __call__(self, imgs):\n        # the buffer is not used\n        if self.buffer_size == 0:\n            return imgs\n        \n        return_imgs = []\n        for img in imgs:\n            img = img.unsqueeze(dim=0)\n            \n            # fill buffer to maximum capacity\n            if self.curr_cap < self.buffer_size:\n                self.curr_cap += 1\n                self.buffer.append(img)\n                return_imgs.append(img)\n            else:\n                p = np.random.uniform(low=0., high=1.)\n                \n                # swap images between input and buffer with probability 0.5\n                if p > 0.5:\n                    idx = np.random.randint(low=0, high=self.buffer_size)\n                    tmp = self.buffer[idx].clone()\n                    self.buffer[idx] = img\n                    return_imgs.append(tmp)\n                else:\n                    return_imgs.append(img)\n        return torch.cat(return_imgs, dim=0)","metadata":{"execution":{"iopub.status.busy":"2023-12-14T04:56:31.017329Z","iopub.execute_input":"2023-12-14T04:56:31.017608Z","iopub.status.idle":"2023-12-14T04:56:31.030715Z","shell.execute_reply.started":"2023-12-14T04:56:31.017585Z","shell.execute_reply":"2023-12-14T04:56:31.029868Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**CycleGan**\n\nWith the generator and discriminator architectures defined, we can now build CycleGAN, which consists of two generators and two discriminators:\n\n* Generator for photo-to-Monet translation (gen_PM).\n* Generator for Monet-to-photo translation (gen_MP).\n* Discriminator for Monet paintings (disc_M).\n* Discriminator for photos (disc_P).\nWe use the Adam optimizer for model training. To optimize the parameters, we need to define the loss functions:\n\n* **Discriminator loss**:- For real images fed into the discriminator, the output matrix is compared against a matrix of 1s using the mean squared error. For fake images, the output matrix is compared against a matrix of 0s. This suggests that to minimize loss, the perfect discriminator outputs a matrix of 1s for real images and a matrix of 0s for fake images.\n* **Generator loss:-** This is composed of three different loss functions below.\n     *          *Adversarial loss* :-Fake images are fed into the discriminator and the output matrix is compared against a matrix of 1s using the mean squared error. To minimize loss, the generator needs to 'fool' the discriminator into thinking that the fake images are real and output a matrix of 1s.\n    *         *Identity loss*:- When a Monet painting is fed into the photo-to-Monet generator, we should get back the same Monet painting because nothing needs to be transformed. The same applies for photos fed into the Monet-to-photo generator. To encourage identity mapping, the difference in pixel values between the input image and generated image is measured using the l1 loss.\n    *         *Cycle loss* When a Monet painting is fed into the Monet-to-photo generator, and the generated image is fed back into the photo-to-Monet generator, it should transform back into the original Monet painting. The same applies for photos passed to the two generators to get back the original photos. To preserve information throughout this cycle, the l1 loss is used to measure the difference between the original image and the reconstructed image.","metadata":{}},{"cell_type":"markdown","source":"**Building the CycleGAN model**","metadata":{}},{"cell_type":"code","source":"MODEL_CONFIG = {\n    # the type of generator, and the number of residual blocks if ResNet generator is used\n    \"gen_name\": \"unet\", # types: 'unet', 'resnet'\n    \"num_resblocks\": 6,\n    # the number of filters in the first layer for the generators and discriminators\n    \"hid_channels\": 64,\n    # using DeepSpeed's FusedAdam (currently GPU only) is slightly faster\n    \"optimizer\": ds.ops.adam.FusedAdam if torch.cuda.is_available() else torch.optim.Adam,\n    # the learning rate and beta parameters for the Adam optimizer\n    \"lr\": 2e-4,\n    \"betas\": (0.5, 0.999),\n    # the weights used in the identity loss and cycle loss\n    \"lambda_idt\": 0.5,\n    \"lambda_cycle\": (10, 10), # (MPM direction, PMP direction)\n    # the size of the buffer that stores previously generated images\n    \"buffer_size\": 100,\n    # the number of epochs for training\n    \"num_epochs\": 18 if not DEBUG else 2,\n    # the number of epochs before starting the learning rate decay\n    \"decay_epochs\": 5 if not DEBUG else 1,\n}","metadata":{"execution":{"iopub.status.busy":"2023-12-14T04:56:31.031944Z","iopub.execute_input":"2023-12-14T04:56:31.032262Z","iopub.status.idle":"2023-12-14T04:56:31.041193Z","shell.execute_reply.started":"2023-12-14T04:56:31.032233Z","shell.execute_reply":"2023-12-14T04:56:31.040333Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CycleGAN(L.LightningModule):\n    def __init__(\n        self,\n        gen_name,\n        num_resblocks,\n        hid_channels,\n        optimizer,\n        lr,\n        betas,\n        lambda_idt,\n        lambda_cycle,\n        buffer_size,\n        num_epochs,\n        decay_epochs,\n    ):\n        super().__init__()\n        self.save_hyperparameters(ignore=[\"optimizer\"])\n        self.optimizer = optimizer\n        self.automatic_optimization = False\n        \n        # define generators and discriminators\n        self.gen_PM = get_gen(gen_name, hid_channels, num_resblocks)\n        self.gen_MP = get_gen(gen_name, hid_channels, num_resblocks)\n        self.disc_M = Discriminator(hid_channels)\n        self.disc_P = Discriminator(hid_channels)\n        \n        # initialize buffers to store fake images\n        self.buffer_fake_M = ImageBuffer(buffer_size)\n        self.buffer_fake_P = ImageBuffer(buffer_size)\n        \n    def forward(self, img):\n        return self.gen_PM(img)   \n            \n    def init_weights(self):\n        def init_fn(m):\n             if hasattr(m, 'weight') and isinstance(m.weight, nn.Parameter):\n                nn.init.normal_(m.weight, 0.0, 0.02)\n                if m.bias is not None:\n                    nn.init.constant_(m.bias, 0.0)\n        \n        for net in [self.gen_PM, self.gen_MP, self.disc_M, self.disc_P]:\n            net.apply(init_fn)\n        \n    def setup(self, stage):\n        if stage == \"fit\":\n            self.init_weights()\n            print(\"Model initialized.\")\n            \n    def get_lr_scheduler(self, optimizer):\n        def lr_lambda(epoch):\n            len_decay_phase = self.hparams.num_epochs - self.hparams.decay_epochs + 1.0\n            curr_decay_step = max(0, epoch - self.hparams.decay_epochs + 1.0)\n            val = 1.0 - curr_decay_step / len_decay_phase\n            return max(0.0, val)\n        \n        return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lr_lambda)\n    \n    def configure_optimizers(self):\n        opt_config = {\n            \"lr\": self.hparams.lr,\n            \"betas\": self.hparams.betas,\n        }\n        opt_gen = self.optimizer(\n            list(self.gen_PM.parameters()) + list(self.gen_MP.parameters()),\n            **opt_config,\n        )\n        opt_disc = self.optimizer(\n            list(self.disc_M.parameters()) + list(self.disc_P.parameters()),\n            **opt_config,\n        )\n        optimizers = [opt_gen, opt_disc]\n        schedulers = [self.get_lr_scheduler(opt) for opt in optimizers]\n        return optimizers, schedulers\n        \n    def adv_criterion(self, y_hat, y):\n        return F.mse_loss(y_hat, y)\n    \n    def recon_criterion(self, y_hat, y):\n        return F.l1_loss(y_hat, y)\n    \n    def get_adv_loss(self, fake, disc):\n        fake_hat = disc(fake)\n        real_labels = torch.ones_like(fake_hat)\n        adv_loss = self.adv_criterion(fake_hat, real_labels)\n        return adv_loss\n    \n    def get_idt_loss(self, real, idt, lambda_cycle):\n        idt_loss = self.recon_criterion(idt, real)\n        return self.hparams.lambda_idt * lambda_cycle * idt_loss\n    \n    def get_cycle_loss(self, real, recon, lambda_cycle):\n        cycle_loss = self.recon_criterion(recon, real)\n        return lambda_cycle * cycle_loss\n    \n    def get_gen_loss(self):\n        # calculate adversarial loss\n        adv_loss_PM = self.get_adv_loss(self.fake_M, self.disc_M)\n        adv_loss_MP = self.get_adv_loss(self.fake_P, self.disc_P)\n        total_adv_loss = adv_loss_PM + adv_loss_MP\n        \n        # calculate identity loss\n        lambda_cycle = self.hparams.lambda_cycle\n        idt_loss_MM = self.get_idt_loss(self.real_M, self.idt_M, lambda_cycle[0])\n        idt_loss_PP = self.get_idt_loss(self.real_P, self.idt_P, lambda_cycle[1])\n        total_idt_loss = idt_loss_MM + idt_loss_PP\n        \n        # calculate cycle loss\n        cycle_loss_MPM = self.get_cycle_loss(self.real_M, self.recon_M, lambda_cycle[0])\n        cycle_loss_PMP = self.get_cycle_loss(self.real_P, self.recon_P, lambda_cycle[1])\n        total_cycle_loss = cycle_loss_MPM + cycle_loss_PMP\n        \n        # combine losses\n        gen_loss = total_adv_loss + total_idt_loss + total_cycle_loss\n        return gen_loss\n    \n    def get_disc_loss(self, real, fake, disc):\n        # calculate loss on real images\n        real_hat = disc(real)\n        real_labels = torch.ones_like(real_hat)\n        real_loss = self.adv_criterion(real_hat, real_labels)\n        \n        # calculate loss on fake images\n        fake_hat = disc(fake.detach())\n        fake_labels = torch.zeros_like(fake_hat)\n        fake_loss = self.adv_criterion(fake_hat, fake_labels)\n        \n        # combine losses\n        disc_loss = (fake_loss + real_loss) * 0.5\n        return disc_loss\n    \n    def get_disc_loss_M(self):\n        fake_M = self.buffer_fake_M(self.fake_M)\n        return self.get_disc_loss(self.real_M, fake_M, self.disc_M)\n    \n    def get_disc_loss_P(self):\n        fake_P = self.buffer_fake_P(self.fake_P)\n        return self.get_disc_loss(self.real_P, fake_P, self.disc_P)\n    \n    def training_step(self, batch, batch_idx):\n        self.real_M = batch[\"monet\"]\n        self.real_P = batch[\"photo\"]\n        opt_gen, opt_disc = self.optimizers()\n\n        # generate fake images\n        self.fake_M = self.gen_PM(self.real_P)\n        self.fake_P = self.gen_MP(self.real_M)\n        \n        # generate identity images\n        self.idt_M = self.gen_PM(self.real_M)\n        self.idt_P = self.gen_MP(self.real_P)\n        \n        # reconstruct images\n        self.recon_M = self.gen_PM(self.fake_P)\n        self.recon_P = self.gen_MP(self.fake_M)\n    \n        # train generators\n        self.toggle_optimizer(opt_gen)\n        gen_loss = self.get_gen_loss()        \n        opt_gen.zero_grad()\n        self.manual_backward(gen_loss)\n        opt_gen.step()\n        self.untoggle_optimizer(opt_gen)\n        \n        # train discriminators\n        self.toggle_optimizer(opt_disc)\n        disc_loss_M = self.get_disc_loss_M()\n        disc_loss_P = self.get_disc_loss_P()\n        opt_disc.zero_grad()\n        self.manual_backward(disc_loss_M)\n        self.manual_backward(disc_loss_P)\n        opt_disc.step()\n        self.untoggle_optimizer(opt_disc)\n        \n        # record training losses\n        metrics = {\n            \"gen_loss\": gen_loss,\n            \"disc_loss_M\": disc_loss_M,\n            \"disc_loss_P\": disc_loss_P,\n        }\n        self.log_dict(metrics, on_step=False, on_epoch=True, prog_bar=True)\n        \n    def validation_step(self, batch, batch_idx):\n        self.display_results(batch, batch_idx, \"validate\")\n    \n    def test_step(self, batch, batch_idx):\n        self.display_results(batch, batch_idx, \"test\")\n        \n    def predict_step(self, batch, batch_idx):\n        return self(batch)\n    \n    def display_results(self, batch, batch_idx, stage):\n        real_P = batch\n        fake_M = self(real_P)\n        \n        if stage == \"validate\":\n            title = f\"Epoch {self.current_epoch+1}: Photo-to-Monet Translation\"\n        else:\n            title = f\"Sample {batch_idx+1}: Photo-to-Monet Translation\"\n\n        show_img(\n            torch.cat([real_P, fake_M], dim=0),\n            nrow=len(real_P),\n            title=title,\n        )\n    \n    def on_train_epoch_start(self):\n        # record learning rates\n        curr_lr = self.lr_schedulers()[0].get_last_lr()[0]\n        self.log(\"lr\", curr_lr, on_step=False, on_epoch=True, prog_bar=True)\n        \n    def on_train_epoch_end(self):\n        # update learning rates\n        for sch in self.lr_schedulers():\n            sch.step()\n        \n        # print current state of epoch\n        logged_values = self.trainer.progress_bar_metrics\n        print(\n            f\"Epoch {self.current_epoch+1}\",\n            *[f\"{k}: {v:.5f}\" for k, v in logged_values.items()],\n            sep=\" - \",\n        )\n        \n    def on_train_end(self):\n        print(\"Training ended.\")\n        \n    def on_predict_epoch_end(self):\n        predictions = self.trainer.predict_loop.predictions\n        num_batches = len(predictions)\n        batch_size = predictions[0].shape[0]\n        last_batch_diff = batch_size - predictions[-1].shape[0]\n        print(f\"Number of images generated: {num_batches*batch_size-last_batch_diff}.\")","metadata":{"execution":{"iopub.status.busy":"2023-12-14T04:56:31.042634Z","iopub.execute_input":"2023-12-14T04:56:31.042978Z","iopub.status.idle":"2023-12-14T04:56:31.079643Z","shell.execute_reply.started":"2023-12-14T04:56:31.042947Z","shell.execute_reply":"2023-12-14T04:56:31.078668Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model Training ","metadata":{}},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"TRAIN_CONFIG = {\n    \"accelerator\": \"gpu\" if torch.cuda.is_available() else \"cpu\",\n    \n    # train on 16-bit precision\n    \"precision\": \"16-mixed\" if torch.cuda.is_available() else 32,\n    \n    # train on single GPU\n    \"devices\": 1,\n    \n    # save checkpoint only for last epoch by default\n    \"enable_checkpointing\": True,\n    \n    # disable logging for simplicity\n    \"logger\": False,\n    \n    # the number of epochs for training (we limit the number of train/predict batches during debugging)\n    \"max_epochs\": MODEL_CONFIG[\"num_epochs\"],\n    \"limit_train_batches\": 1.0 if not DEBUG else 2,\n    \"limit_predict_batches\": 1.0 if not DEBUG else 5,\n    \n    # the maximum amount of time for training, in case we exceed run-time of 5 hours\n    \"max_time\": {\"hours\": 4, \"minutes\": 55},\n    \n    # use a small subset of photos for validation/testing (we limit here for flexibility)\n    \"limit_val_batches\": 1,\n    \"limit_test_batches\": 5,\n    \n    # disable sanity check before starting the training routine\n    \"num_sanity_val_steps\": 0,\n    \n    # the frequency to visualize the progress of adding Monet style\n    \"check_val_every_n_epoch\": 2 if not DEBUG else 1,\n}","metadata":{"execution":{"iopub.status.busy":"2023-12-14T04:56:31.080821Z","iopub.execute_input":"2023-12-14T04:56:31.081176Z","iopub.status.idle":"2023-12-14T04:56:31.092976Z","shell.execute_reply.started":"2023-12-14T04:56:31.081143Z","shell.execute_reply":"2023-12-14T04:56:31.092188Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"dm = CustomDataModule(**DM_CONFIG)\nmodel = CycleGAN(**MODEL_CONFIG)\ntrainer = L.Trainer(**TRAIN_CONFIG)\ntrainer.fit(model, datamodule=dm)","metadata":{"execution":{"iopub.status.busy":"2023-12-14T04:56:31.093992Z","iopub.execute_input":"2023-12-14T04:56:31.094265Z","iopub.status.idle":"2023-12-14T05:02:57.654399Z","shell.execute_reply.started":"2023-12-14T04:56:31.094242Z","shell.execute_reply":"2023-12-14T05:02:57.653160Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Testing the model on other sample photos**","metadata":{}},{"cell_type":"code","source":"_ = trainer.test(model, datamodule=dm)","metadata":{"execution":{"iopub.status.busy":"2023-12-14T05:02:57.656139Z","iopub.execute_input":"2023-12-14T05:02:57.656570Z","iopub.status.idle":"2023-12-14T05:03:00.794496Z","shell.execute_reply.started":"2023-12-14T05:02:57.656528Z","shell.execute_reply":"2023-12-14T05:03:00.793491Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Submission","metadata":{}},{"cell_type":"markdown","source":"Computing the predictions can be done by running the predict method to generate the Monet-style images given the input photos.","metadata":{}},{"cell_type":"code","source":"predictions = trainer.predict(model, datamodule=dm)","metadata":{"execution":{"iopub.status.busy":"2023-12-14T05:03:00.795928Z","iopub.execute_input":"2023-12-14T05:03:00.796229Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Saving the generated images**","metadata":{}},{"cell_type":"code","source":"os.makedirs(\"../images\", exist_ok=True)\nidx = 0\nfor tensor in predictions:\n    for monet in tensor:\n        save_image(\n            monet.float().squeeze() * 0.5 + 0.5, \n            fp=f\"../images/{idx}.jpg\",\n        )\n        idx += 1\n\nshutil.make_archive(\"/kaggle/working/images\", \"zip\", \"/kaggle/images\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}