{"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":"markdown","source":"# Photo to Monet — CycleGAN⚡PyTorch Lightning","metadata":{"papermill":{"duration":0.012975,"end_time":"2023-03-19T01:57:31.956786","exception":false,"start_time":"2023-03-19T01:57:31.943811","status":"completed"},"tags":[]}},{"cell_type":"markdown","source":"This notebook aims to implement CycleGAN, with the model architecture being adapted from the [tutorial](https://www.kaggle.com/code/amyjang/monet-cyclegan-tutorial/notebook) available. We attempt to convert the code to PyTorch Lightning here. A useful reference for basic GAN in PyTorch Lightning can be found in this [documentation](https://pytorch-lightning.readthedocs.io/en/stable/notebooks/lightning_examples/basic-gan.html). More details regarding CycleGAN can be found in the original [paper](https://arxiv.org/abs/1703.10593) and [code](https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix).","metadata":{"papermill":{"duration":0.010961,"end_time":"2023-03-19T01:57:31.979038","exception":false,"start_time":"2023-03-19T01:57:31.968077","status":"completed"},"tags":[]}},{"cell_type":"code","source":"import glob\nimport os\nimport shutil\n\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.trainer.supporters 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\n_ = L.seed_everything(0, workers=True)","metadata":{"_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","papermill":{"duration":12.994572,"end_time":"2023-03-19T01:57:44.984673","exception":false,"start_time":"2023-03-19T01:57:31.990101","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-04-13T20:16:11.944424Z","iopub.execute_input":"2023-04-13T20:16:11.944932Z","iopub.status.idle":"2023-04-13T20:16:18.318276Z","shell.execute_reply.started":"2023-04-13T20:16:11.944886Z","shell.execute_reply":"2023-04-13T20:16:18.316961Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"---","metadata":{"papermill":{"duration":0.01099,"end_time":"2023-03-19T01:57:45.007205","exception":false,"start_time":"2023-03-19T01:57:44.996215","status":"completed"},"tags":[]}},{"cell_type":"markdown","source":"# 1. Data Preprocessing","metadata":{"papermill":{"duration":0.011194,"end_time":"2023-03-19T01:57:45.030579","exception":false,"start_time":"2023-03-19T01:57:45.019385","status":"completed"},"tags":[]}},{"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=(18, 8))\n    plt.imshow(img_grid)\n    plt.axis(\"off\")\n    plt.title(title)\n    plt.show()","metadata":{"_kg_hide-input":false,"papermill":{"duration":0.020996,"end_time":"2023-03-19T01:57:45.062727","exception":false,"start_time":"2023-03-19T01:57:45.041731","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-04-13T20:16:18.320864Z","iopub.execute_input":"2023-04-13T20:16:18.321162Z","iopub.status.idle":"2023-04-13T20:16:18.326892Z","shell.execute_reply.started":"2023-04-13T20:16:18.321125Z","shell.execute_reply":"2023-04-13T20:16:18.325853Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Augmenting the images.","metadata":{"papermill":{"duration":0.010988,"end_time":"2023-03-19T01:57:45.084838","exception":false,"start_time":"2023-03-19T01:57:45.07385","status":"completed"},"tags":[]}},{"cell_type":"markdown","source":"Before 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 in every epoch, especially when we only have 300 Monet paintings. We look at some basic image transformations:\n* Cropping a random portion of the image and resizing it to the original size of 256 with `RandomResizedCrop` to obtain slightly different images.\n* Randomly flipping the images using `RandomHorizontalFlip`.\n\nOther possible transformations can be found [here](https://pytorch.org/vision/stable/transforms.html). These transformations are only needed during model training/fitting, and we specify this using the `stage` argument. Finally, the images are scaled down for better convergence. The above transformations are chosen to avoid large changes to the distributions of the real images, that is the transformed images are still interpretable (as compared to the case of flipping MNIST handwritten digits). ","metadata":{"papermill":{"duration":0.010992,"end_time":"2023-03-19T01:57:45.1071","exception":false,"start_time":"2023-03-19T01:57:45.096108","status":"completed"},"tags":[]}},{"cell_type":"code","source":"class CustomTransform(object):\n    def __init__(self, img_dim=256):\n        self.transform = T.Compose([\n            T.Resize((img_dim, img_dim)),\n            T.RandomResizedCrop((img_dim, img_dim),\n                                scale=(0.7, 0.9), ratio=(0.9, 1.1)),\n            T.RandomHorizontalFlip(p=0.5),\n        ])\n        \n    def __call__(self, img, stage=\"fit\"):\n        if stage == \"fit\" or stage == \"predict\":\n            img = self.transform(img)\n        return img/127.5 - 1","metadata":{"papermill":{"duration":0.021699,"end_time":"2023-03-19T01:57:45.140105","exception":false,"start_time":"2023-03-19T01:57:45.118406","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-04-13T20:16:18.328342Z","iopub.execute_input":"2023-04-13T20:16:18.328943Z","iopub.status.idle":"2023-04-13T20:16:18.337739Z","shell.execute_reply.started":"2023-04-13T20:16:18.328906Z","shell.execute_reply":"2023-04-13T20:16:18.336753Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Storing the datasets.","metadata":{"papermill":{"duration":0.011074,"end_time":"2023-03-19T01:57:45.162356","exception":false,"start_time":"2023-03-19T01:57:45.151282","status":"completed"},"tags":[]}},{"cell_type":"markdown","source":"To load and store the datasets, we define a custom `Dataset` involving three main methods: \n* `__init__` to initialize the dataset.\n* `__len__` to retrieve the size of the dataset.\n* `__getitem__` to get the i-th sample of images after performing the transformations described above.\n\nSimilarly, we define the `stage` argument to differentiate between the training dataset and prediction dataset when performing the transformations. Different instances of `CustomDataset` will be used to retrieve the photos and Monet paintings separately. We look at combining them while iterating through the datasets later.","metadata":{"papermill":{"duration":0.011056,"end_time":"2023-03-19T01:57:45.18459","exception":false,"start_time":"2023-03-19T01:57:45.173534","status":"completed"},"tags":[]}},{"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)\n        return self.transform(img, stage=self.stage)","metadata":{"papermill":{"duration":0.025818,"end_time":"2023-03-19T01:57:45.221633","exception":false,"start_time":"2023-03-19T01:57:45.195815","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-04-13T20:16:18.340659Z","iopub.execute_input":"2023-04-13T20:16:18.341292Z","iopub.status.idle":"2023-04-13T20:16:18.347782Z","shell.execute_reply.started":"2023-04-13T20:16:18.341255Z","shell.execute_reply":"2023-04-13T20:16:18.346821Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Iterating through the datasets.","metadata":{"papermill":{"duration":0.011352,"end_time":"2023-03-19T01:57:45.244912","exception":false,"start_time":"2023-03-19T01:57:45.23356","status":"completed"},"tags":[]}},{"cell_type":"markdown","source":"To 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. We specify the sampling mode using `mode=\"max_size_cycle\"` to stop after one complete pass of the larger dataset of photos while cycling through the smaller dataset of Monet paintings. Other modes can be found [here](https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.utilities.combined_loader.html). In contrast, we fix the prediction dataset to only contain photos for simplicity since we are trying to generate Monet-style images in this notebook. As such, we do not combine any dataloaders for the prediction dataset.\n\nTo organize all the steps described above for processing data, we define a custom `LightningDataModule`. A datamodule involves many methods, but we are mainly concerned with:\n* `setup` to create the datasets and apply the corresponding transformations defined above.\n* `train_dataloader` to generate the dataloader for the training dataset.\n* `predict_dataloader` to generate the dataloader for the prediction dataset.\n\nOther possible methods can be found [here](https://pytorch-lightning.readthedocs.io/en/stable/data/datamodule.html). We define the following parameters:\n* `MONET_DIR` and `PHOTO_DIR` — the directories where the Monet paintings and photos are loaded from respectively.\n* `batch_size` — the number of samples in each training/prediction batch.\n* `num_workers` — the number of subprocesses (excluding the main process) used for data loading.\n* `pin_memory` — to enable faster data transfer to GPU during training.","metadata":{"papermill":{"duration":0.011068,"end_time":"2023-03-19T01:57:45.267418","exception":false,"start_time":"2023-03-19T01:57:45.25635","status":"completed"},"tags":[]}},{"cell_type":"code","source":"MONET_DIR = \"/kaggle/input/gan-getting-started/monet_jpg/*.jpg\"\nPHOTO_DIR = \"/kaggle/input/gan-getting-started/photo_jpg/*.jpg\"\nPREDICT_DIR = \"/kaggle/input/Homework/CA3/Part 2-4/*.jpg\"\nLOADER_CONFIG = {\n    \"batch_size\": 8,\n    \"num_workers\": os.cpu_count(),\n    \"pin_memory\": torch.cuda.is_available(),\n}\ntransform = CustomTransform()","metadata":{"papermill":{"duration":0.019339,"end_time":"2023-03-19T01:57:45.297921","exception":false,"start_time":"2023-03-19T01:57:45.278582","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-04-13T20:16:18.349036Z","iopub.execute_input":"2023-04-13T20:16:18.349989Z","iopub.status.idle":"2023-04-13T20:16:18.407351Z","shell.execute_reply.started":"2023-04-13T20:16:18.349948Z","shell.execute_reply":"2023-04-13T20:16:18.406309Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CustomDataModule(L.LightningDataModule):\n    def __init__(\n        self, \n        monet_dir=MONET_DIR,\n        photo_dir=PHOTO_DIR,        \n        predict_dir=PREDICT_DIR,\n        loader_config=LOADER_CONFIG,\n        transform=transform,\n        mode=\"max_size_cycle\",\n    ):\n        super().__init__()\n        self.monet_filenames = sorted(glob.glob(monet_dir))\n        self.photo_filenames = sorted(glob.glob(photo_dir))\n        self.predict_filenames = sorted(glob.glob(predict_dir))\n        self.loader_config = loader_config\n        self.transform = transform\n        self.mode = mode\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        elif stage == \"predict\":\n            self.predict = CustomDataset(self.predict_filenames, self.transform, stage)\n            \n    def train_dataloader(self):\n        loader_monet = DataLoader(\n            self.train_monet, \n            shuffle=True, \n            drop_last=True, \n            **self.loader_config,\n        )\n        loader_photo = DataLoader(\n            self.train_photo, \n            shuffle=True,\n            drop_last=True,\n            **self.loader_config,\n        )\n        loaders = {\"monet\": loader_monet, \"photo\": loader_photo}\n        return CombinedLoader(loaders, mode=self.mode)\n    \n    def predict_dataloader(self):\n        return DataLoader(self.predict, shuffle=False, **self.loader_config)","metadata":{"papermill":{"duration":0.022601,"end_time":"2023-03-19T01:57:45.331838","exception":false,"start_time":"2023-03-19T01:57:45.309237","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-04-13T20:16:18.411063Z","iopub.execute_input":"2023-04-13T20:16:18.411334Z","iopub.status.idle":"2023-04-13T20:16:18.421698Z","shell.execute_reply.started":"2023-04-13T20:16:18.411307Z","shell.execute_reply":"2023-04-13T20:16:18.420312Z"},"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":{"papermill":{"duration":0.011244,"end_time":"2023-03-19T01:57:45.354335","exception":false,"start_time":"2023-03-19T01:57:45.343091","status":"completed"},"tags":[]}},{"cell_type":"code","source":"SAMPLE_SIZE=5\ndm_sample = CustomDataModule(loader_config={\"batch_size\": SAMPLE_SIZE})\n\ndm_sample.setup(\"fit\")\ntrain_loader = dm_sample.train_dataloader()\nmonet_samples = next(iter(train_loader))[\"monet\"]\n\ndm_sample.setup(\"predict\")\npredict_loader = dm_sample.predict_dataloader()\nphoto_samples = next(iter(predict_loader))\n\nshow_img(monet_samples, nrow=SAMPLE_SIZE, title=\"Samples of Monet Paintings\")\nshow_img(photo_samples, nrow=SAMPLE_SIZE, title=\"Samples of Photos\")","metadata":{"papermill":{"duration":1.301484,"end_time":"2023-03-19T01:57:46.667097","exception":false,"start_time":"2023-03-19T01:57:45.365613","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-04-13T20:16:18.423165Z","iopub.execute_input":"2023-04-13T20:16:18.423965Z","iopub.status.idle":"2023-04-13T20:16:20.239513Z","shell.execute_reply.started":"2023-04-13T20:16:18.423927Z","shell.execute_reply":"2023-04-13T20:16:20.237663Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"---","metadata":{"papermill":{"duration":0.035229,"end_time":"2023-03-19T01:57:46.737485","exception":false,"start_time":"2023-03-19T01:57:46.702256","status":"completed"},"tags":[]}},{"cell_type":"markdown","source":"# 2. Building CycleGAN Architecture","metadata":{"papermill":{"duration":0.032736,"end_time":"2023-03-19T01:57:46.80351","exception":false,"start_time":"2023-03-19T01:57:46.770774","status":"completed"},"tags":[]}},{"cell_type":"markdown","source":"### Generator.","metadata":{}},{"cell_type":"markdown","source":"<img src=\"https://lh5.googleusercontent.com/9kNO6hxYJmpcfG5bOjnDazieeLC7Q8jZJi3gTtnJelbkOUL7Xz9e-3F_SNuxPpo4fZ4=w2400\" width=\"600\"/>\n    \n_Example of the U-Net architecture [[source](https://paperswithcode.com/method/u-net)]._","metadata":{"papermill":{"duration":0.038219,"end_time":"2023-03-19T01:57:46.874277","exception":false,"start_time":"2023-03-19T01:57:46.836058","status":"completed"},"tags":[]}},{"cell_type":"markdown","source":"We use a U-Net architecture for the CycleGAN generator. U-Net is a network which consists of downsampling blocks and upsampling blocks with skip connections, giving it the U-shaped architecture.","metadata":{"papermill":{"duration":0.032886,"end_time":"2023-03-19T01:57:46.94176","exception":false,"start_time":"2023-03-19T01:57:46.908874","status":"completed"},"tags":[]}},{"cell_type":"markdown","source":"### Downsampling blocks.","metadata":{"papermill":{"duration":0.032774,"end_time":"2023-03-19T01:57:47.00759","exception":false,"start_time":"2023-03-19T01:57:46.974816","status":"completed"},"tags":[]}},{"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":{"papermill":{"duration":0.032597,"end_time":"2023-03-19T01:57:47.073128","exception":false,"start_time":"2023-03-19T01:57:47.040531","status":"completed"},"tags":[]}},{"cell_type":"code","source":"class Downsampling(nn.Module):\n    def __init__(\n        self,\n        in_channels,\n        out_channels,\n        norm=True,\n        kernel_size=4,\n        stride=2,\n        padding=1,\n    ):\n        super().__init__()\n        self.block = nn.Sequential(\n            nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size,\n                      stride=stride, padding=padding, bias=False)\n        )\n        if norm:\n            self.block.append(nn.InstanceNorm2d(out_channels, affine=True))\n        self.block.append(nn.LeakyReLU(0.2))\n        \n    def forward(self, x):\n        return self.block(x)","metadata":{"papermill":{"duration":0.044778,"end_time":"2023-03-19T01:57:47.150705","exception":false,"start_time":"2023-03-19T01:57:47.105927","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-04-13T20:16:20.241067Z","iopub.execute_input":"2023-04-13T20:16:20.241808Z","iopub.status.idle":"2023-04-13T20:16:20.250089Z","shell.execute_reply.started":"2023-04-13T20:16:20.241756Z","shell.execute_reply":"2023-04-13T20:16:20.249038Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Upsampling blocks.","metadata":{"papermill":{"duration":0.032956,"end_time":"2023-03-19T01:57:47.216301","exception":false,"start_time":"2023-03-19T01:57:47.183345","status":"completed"},"tags":[]}},{"cell_type":"markdown","source":"On the other hand, the upsampling blocks contain transposed convolution layers, which combine the learned features to output an image with the original size 256.","metadata":{"papermill":{"duration":0.032817,"end_time":"2023-03-19T01:57:47.282273","exception":false,"start_time":"2023-03-19T01:57:47.249456","status":"completed"},"tags":[]}},{"cell_type":"code","source":"class Upsampling(nn.Module):\n    def __init__(\n        self,\n        in_channels,\n        out_channels,\n        dropout=False,\n        kernel_size=4,\n        stride=2,\n        padding=1,\n        output_padding=0,\n    ):\n        super().__init__()\n        self.block = nn.Sequential(\n            nn.ConvTranspose2d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, \n                               padding=padding, output_padding=output_padding, bias=False),\n            nn.InstanceNorm2d(out_channels, affine=True),\n        )\n        if dropout:\n            self.block.append(nn.Dropout(0.5))\n        self.block.append(nn.ReLU())\n        \n    def forward(self, x):\n        return self.block(x)","metadata":{"papermill":{"duration":0.044362,"end_time":"2023-03-19T01:57:47.359186","exception":false,"start_time":"2023-03-19T01:57:47.314824","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-04-13T20:16:20.251563Z","iopub.execute_input":"2023-04-13T20:16:20.252186Z","iopub.status.idle":"2023-04-13T20:16:20.260987Z","shell.execute_reply.started":"2023-04-13T20:16:20.252144Z","shell.execute_reply":"2023-04-13T20:16:20.260143Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Building the generator.","metadata":{}},{"cell_type":"markdown","source":"With the building blocks defined, we can now build our CycleGAN generator. 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. For reference, the output size of each block is commented below.","metadata":{}},{"cell_type":"code","source":"class Generator(nn.Module):\n    def __init__(self, in_channels, out_channels, hid_channels):\n        super().__init__()\n        self.downsampling_blocks = 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), # 512x16x16\n            Downsampling(hid_channels*8, hid_channels*8), # 512x8x8\n            Downsampling(hid_channels*8, hid_channels*8), # 512x4x4\n            Downsampling(hid_channels*8, hid_channels*8), # 512x2x2\n            Downsampling(hid_channels*8, hid_channels*8, norm=False), # 512x1x1\n        )\n        self.upsampling_blocks = nn.Sequential(\n            Upsampling(hid_channels*8, hid_channels*8, dropout=True), # (512+512)x2x2\n            Upsampling(hid_channels*16, hid_channels*8, dropout=True), # (512+512)x4x4\n            Upsampling(hid_channels*16, hid_channels*8, dropout=True), # (512+512)x8x8\n            Upsampling(hid_channels*16, hid_channels*8), # (512+512)x16x16\n            Upsampling(hid_channels*16, hid_channels*4), # (256+256)x32x32\n            Upsampling(hid_channels*8, hid_channels*2), # (128+128)x64x64\n            Upsampling(hid_channels*4, hid_channels), # (64+64)x128x128\n        )\n        self.feature_block = nn.Sequential(\n            nn.ConvTranspose2d(hid_channels*2, out_channels,\n                               kernel_size=4, stride=2, padding=1), # 3x256x256\n            nn.Tanh(),\n        )\n        \n    def forward(self, x):\n        skips = []\n        for down in self.downsampling_blocks:\n            x = down(x)\n            skips.append(x)\n        skips = reversed(skips[:-1])\n        \n        for up, skip in zip(self.upsampling_blocks, skips):\n            x = up(x)\n            x = torch.cat([x, skip], dim=1)\n        return self.feature_block(x)","metadata":{"papermill":{"duration":0.047175,"end_time":"2023-03-19T01:57:48.061105","exception":false,"start_time":"2023-03-19T01:57:48.01393","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-04-13T20:16:20.265667Z","iopub.execute_input":"2023-04-13T20:16:20.266513Z","iopub.status.idle":"2023-04-13T20:16:20.279889Z","shell.execute_reply.started":"2023-04-13T20:16:20.266475Z","shell.execute_reply":"2023-04-13T20:16:20.278907Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Discriminator.","metadata":{"papermill":{"duration":0.033167,"end_time":"2023-03-19T01:57:48.132136","exception":false,"start_time":"2023-03-19T01:57:48.098969","status":"completed"},"tags":[]}},{"cell_type":"markdown","source":"<div>\n<img src=\"https://lh6.googleusercontent.com/UhJiaTOQWgfHQlWq50IMGBvdkJ3NDggC449cxud8XVlSxUrule8f5LyoLUV8aaYemGw=w2400\" width=\"300\"/>\n</div>\n\n_Diagram of how the PatchGAN discriminator works [[source](https://www.researchgate.net/figure/PatchGAN-discriminator-Each-value-of-the-output-matrix-represents-the-probability-of_fig1_323904616)]._","metadata":{"papermill":{"duration":0.033278,"end_time":"2023-03-19T01:57:48.198341","exception":false,"start_time":"2023-03-19T01:57:48.165063","status":"completed"},"tags":[]}},{"cell_type":"markdown","source":"Unlike 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":{"papermill":{"duration":0.034539,"end_time":"2023-03-19T01:57:48.266305","exception":false,"start_time":"2023-03-19T01:57:48.231766","status":"completed"},"tags":[]}},{"cell_type":"markdown","source":"### Building the discriminator.","metadata":{}},{"cell_type":"markdown","source":"In general, the PatchGAN discriminator consists of a sequence of convolution layers, which can be built using the downsampling blocks defined earlier. For reference, the output size of each block is commented below.","metadata":{}},{"cell_type":"code","source":"class Discriminator(nn.Module):\n    def __init__(self, in_channels, hid_channels):\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":{"papermill":{"duration":0.043903,"end_time":"2023-03-19T01:57:48.343071","exception":false,"start_time":"2023-03-19T01:57:48.299168","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-04-13T20:16:20.281485Z","iopub.execute_input":"2023-04-13T20:16:20.282172Z","iopub.status.idle":"2023-04-13T20:16:20.291961Z","shell.execute_reply.started":"2023-04-13T20:16:20.282126Z","shell.execute_reply":"2023-04-13T20:16:20.290972Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### CycleGAN.","metadata":{"papermill":{"duration":0.044686,"end_time":"2023-03-19T01:57:48.419994","exception":false,"start_time":"2023-03-19T01:57:48.375308","status":"completed"},"tags":[]}},{"cell_type":"markdown","source":"With the generator and discriminator defined, we can now build CycleGAN, which consists of two generators and two discriminators:\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`).\n\nUsing `weights_init` and `reset_parameters` functions, the weights of the layers in the generators and discriminators are initialized using the normal distribution and the biases are initialized to 0s. The Adam optimizer is used for model training. To optimize the parameters, we need to define the loss functions:\n* **Discriminator loss** (`disc_loss`). For real images fed into the discriminator, the output matrix is compared against a matrix of 1s using the binary cross-entropy. 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. To avoid overconfidence in the predictions, we use a technique called one-sided label smoothing, which involves replacing the target matrix of 1s for the real images with values slightly smaller than 1, such as 0.9. This can be controlled by the `smooth` parameter. Label smoothing is done only for real images to reinforce the correct behavior in the generators. More information can be found in this [paper](https://arxiv.org/abs/1701.00160).\n* **Generator loss** (`gen_loss`). This is composed of three different loss functions below.\n  * *Adversarial loss* (`adv_loss`). Fake images are fed into the discriminator and the output matrix is compared against a matrix of 1s using the binary cross-entropy. 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* (`id_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* (`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 cycled image.\n\nFrom the above, the binary cross-entropy and the l1 loss are defined as the adversarial criterion (`adv_criterion`) and the reconstruction criterion (`recon_criterion`) respectively.","metadata":{"papermill":{"duration":0.033641,"end_time":"2023-03-19T01:57:48.487251","exception":false,"start_time":"2023-03-19T01:57:48.45361","status":"completed"},"tags":[]}},{"cell_type":"markdown","source":"### Building the CycleGAN model.","metadata":{}},{"cell_type":"markdown","source":"To organize the code for modeling, we define the above functions within the `LightningModule` class together with the following methods:\n* `__init__` to initialize the two generators, the two discriminators, and other parameters.\n* `forward` to generate Monet-style images given input photos.\n* `configure_optimizers` to define the Adam optimizers.\n* `training_step` to compute the loss functions for the generators and discriminators.\n* `training_epoch_end` to print the average values of the loss functions over the batches per epoch, and visualize the performance of `gen_PM`.\n* `predict_step` to run the `forward` method during prediction.\n\nOther useful methods can be found [here](https://pytorch-lightning.readthedocs.io/en/stable/common/lightning_module.html). The following parameters are set:\n* `IN_CHANNELS` — the number of input channels for the generator and discriminator, which equals 3 since we are working with RGB images.\n* `OUT_CHANNELS` — the number of output channels for the generator, which equals 3 as we are trying to output RGB images.\n* `HID_CHANNELS` — the number of output channels in the first layer for the generator and discriminator.\n* `LR` and `BETAS` — the initial learning rate and beta parameters for the Adam optimizer.\n* `LAMBDA` — the weight used in the identity loss and cycle loss.\n* `DISPLAY_EPOCHS` — the frequency to visualize the performance of `gen_PM`.","metadata":{"papermill":{"duration":0.033641,"end_time":"2023-03-19T01:57:48.487251","exception":false,"start_time":"2023-03-19T01:57:48.45361","status":"completed"},"tags":[]}},{"cell_type":"code","source":"IN_CHANNELS = 3\nOUT_CHANNELS = 3\nHID_CHANNELS = 64\nLR = 2e-4\nBETAS = (0.5, 0.999)\nLAMBDA = 10\nDISPLAY_EPOCHS = 3","metadata":{"papermill":{"duration":0.042883,"end_time":"2023-03-19T01:57:48.564164","exception":false,"start_time":"2023-03-19T01:57:48.521281","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-04-13T20:16:20.293938Z","iopub.execute_input":"2023-04-13T20:16:20.294755Z","iopub.status.idle":"2023-04-13T20:16:20.302305Z","shell.execute_reply.started":"2023-04-13T20:16:20.294715Z","shell.execute_reply":"2023-04-13T20:16:20.301369Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CycleGAN(L.LightningModule):\n    def __init__(\n        self, \n        in_channels=IN_CHANNELS,\n        out_channels=OUT_CHANNELS, \n        hid_channels=HID_CHANNELS,\n        lr=LR,\n        betas=BETAS,\n        lambda_w=LAMBDA,\n        display_epochs=DISPLAY_EPOCHS,\n        photo_samples=photo_samples,\n    ):\n        super().__init__()\n        self.lr = lr\n        self.betas = betas\n        self.lambda_w = lambda_w\n        self.display_epochs = display_epochs\n        self.photo_samples = photo_samples\n        self.loss_history = []\n        self.epoch_count = 0\n        \n        self.gen_PM = Generator(in_channels, out_channels, hid_channels)\n        self.gen_MP = Generator(in_channels, out_channels, hid_channels)\n        self.disc_M = Discriminator(in_channels, hid_channels)\n        self.disc_P = Discriminator(in_channels, hid_channels)\n        self.reset_parameters()\n        \n    def forward(self, z):\n        return self.gen_PM(z)\n    \n    def weights_init(self, m):\n        if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d, nn.InstanceNorm2d)):\n            nn.init.normal_(m.weight, 0., 0.02)\n            if m.bias is not None:\n                nn.init.constant_(m.bias, 0.)\n                \n    def reset_parameters(self):\n        self.gen_PM = self.gen_PM.apply(self.weights_init)\n        self.gen_MP = self.gen_MP.apply(self.weights_init)\n        self.disc_M = self.disc_M.apply(self.weights_init)\n        self.disc_P = self.disc_P.apply(self.weights_init)\n        \n    def adv_criterion(self, y_hat, y):\n        return F.binary_cross_entropy_with_logits(y_hat, y)\n    \n    def recon_criterion(self, y_hat, y):\n        return F.l1_loss(y_hat, y)\n    \n    def adv_loss(self, fake_Y, disc_Y):\n        fake_Y_hat = disc_Y(fake_Y)\n        valid = torch.ones_like(fake_Y_hat)\n        adv_loss_XY = self.adv_criterion(fake_Y_hat, valid)\n        return adv_loss_XY\n    \n    def id_loss(self, real_Y, gen_XY):\n        id_Y = gen_XY(real_Y)\n        id_loss_Y = self.recon_criterion(id_Y, real_Y)\n        return 0.5 * self.lambda_w * id_loss_Y\n    \n    def cycle_loss(self, real_Y, fake_X, gen_XY):\n        cycle_Y = gen_XY(fake_X)\n        cycle_loss_Y = self.recon_criterion(cycle_Y, real_Y)\n        return self.lambda_w * cycle_loss_Y\n        \n    def gen_loss(self, real_X, real_Y, gen_XY, gen_YX, disc_Y):\n        adv_loss_XY = self.adv_loss(gen_XY(real_X), disc_Y)\n\n        id_loss_Y = self.id_loss(real_Y, gen_XY)\n        \n        cycle_loss_Y = self.cycle_loss(real_Y, gen_YX(real_Y), gen_XY)\n        cycle_loss_X = self.cycle_loss(real_X, gen_XY(real_X), gen_YX)\n        total_cycle_loss = cycle_loss_X + cycle_loss_Y\n        \n        gen_loss_XY = adv_loss_XY + id_loss_Y + total_cycle_loss\n        return gen_loss_XY\n    \n    def disc_loss(self, real_Y, fake_Y, disc_Y, smooth=0.9):\n        real_Y_hat = disc_Y(real_Y)\n        valid = torch.full_like(real_Y_hat, fill_value=smooth)\n        real_loss_Y = self.adv_criterion(real_Y_hat, valid)\n        \n        fake_Y_hat = disc_Y(fake_Y.detach())\n        fake = torch.zeros_like(fake_Y_hat)\n        fake_loss_Y = self.adv_criterion(fake_Y_hat, fake)\n        \n        disc_loss_Y = (fake_loss_Y+real_loss_Y) * 0.5\n        return disc_loss_Y\n    \n    def configure_optimizers(self):\n        params = {\n            \"lr\": self.lr,\n            \"betas\": self.betas,\n        }\n        opt_gen_PM = torch.optim.Adam(self.gen_PM.parameters(), **params)\n        opt_gen_MP = torch.optim.Adam(self.gen_MP.parameters(), **params)\n        opt_disc_M = torch.optim.Adam(self.disc_M.parameters(), **params)\n        opt_disc_P = torch.optim.Adam(self.disc_P.parameters(), **params)\n        \n        return [opt_gen_PM, opt_gen_MP, opt_disc_M, opt_disc_P], []\n    \n    def training_step(self, batch, batch_idx, optimizer_idx):\n        real_M = batch[\"monet\"]\n        real_P = batch[\"photo\"]\n        if optimizer_idx == 0:\n            gen_loss_PM = self.gen_loss(real_P, real_M, \n                                        self.gen_PM, self.gen_MP, self.disc_M)\n            return gen_loss_PM\n        \n        if optimizer_idx == 1:\n            gen_loss_MP = self.gen_loss(real_M, real_P,\n                                        self.gen_MP, self.gen_PM, self.disc_P)\n            return gen_loss_MP\n        \n        if optimizer_idx == 2:\n            disc_loss_M = self.disc_loss(real_M, self.gen_PM(real_P), self.disc_M)\n            return disc_loss_M\n        \n        if optimizer_idx == 3:\n            disc_loss_P = self.disc_loss(real_P, self.gen_MP(real_M), self.disc_P)\n            return disc_loss_P\n    \n    def training_epoch_end(self, outputs):\n        self.epoch_count += 1\n        \n        losses = []\n        for j in range(4):\n            loss = np.mean([out[j][\"loss\"].item() for out in outputs])\n            losses.append(loss)\n        self.loss_history.append(losses)\n        print(\n            f\"Epoch {self.epoch_count} -\",\n            f\"gen_loss_PM: {losses[0]:.5f} -\",\n            f\"gen_loss_MP: {losses[1]:.5f} -\",\n            f\"disc_loss_M: {losses[2]:.5f} -\",\n            f\"disc_loss_P: {losses[3]:.5f}\",\n        )\n        \n        if self.epoch_count%self.display_epochs==0 or self.epoch_count==1:\n            gen_monets = self(self.photo_samples.to(self.device)).detach().cpu()\n            show_img(\n                torch.cat([self.photo_samples, gen_monets]),\n                nrow=len(self.photo_samples),\n                title=f\"Epoch {self.epoch_count}: Photo-to-Monet Translation\",\n            )\n            \n    def predict_step(self, batch, batch_idx):\n        return self(batch)\n    \n    def loss_curves(self):\n        labels = [\"gen_loss_PM\", \"gen_loss_MP\", \"disc_loss_M\", \"disc_loss_P\"]\n        titles = [\"Generator Loss Curves\", \"Discriminator Loss Curves\"]\n        num_epochs = len(self.loss_history)\n        plt.figure(figsize=(18, 4.5))\n        for j in range(4):\n            if j%2 == 0:\n                plt.subplot(1, 2, (j//2)+1)\n                plt.title(titles[j//2])\n                plt.ylabel(\"Loss\")\n                plt.xlabel(\"Epoch\")\n            plt.plot(\n                np.arange(1, num_epochs+1),\n                [losses[j] for losses in self.loss_history],\n                label=labels[j],\n            )\n            plt.legend(loc=\"upper right\")","metadata":{"papermill":{"duration":0.064437,"end_time":"2023-03-19T01:57:48.661507","exception":false,"start_time":"2023-03-19T01:57:48.59707","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-04-13T20:16:20.30392Z","iopub.execute_input":"2023-04-13T20:16:20.30464Z","iopub.status.idle":"2023-04-13T20:16:20.333838Z","shell.execute_reply.started":"2023-04-13T20:16:20.304604Z","shell.execute_reply":"2023-04-13T20:16:20.333161Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"---","metadata":{"papermill":{"duration":0.033909,"end_time":"2023-03-19T01:57:48.728836","exception":false,"start_time":"2023-03-19T01:57:48.694927","status":"completed"},"tags":[]}},{"cell_type":"markdown","source":"# 3. Model Training","metadata":{"papermill":{"duration":0.032703,"end_time":"2023-03-19T01:57:48.794464","exception":false,"start_time":"2023-03-19T01:57:48.761761","status":"completed"},"tags":[]}},{"cell_type":"markdown","source":"To start training the model, we use `Trainer` to automatically handle the training loop by running the `fit` method. We set the below parameters in `TRAIN_CONFIG` for training.","metadata":{"papermill":{"duration":0.032684,"end_time":"2023-03-19T01:57:48.859924","exception":false,"start_time":"2023-03-19T01:57:48.82724","status":"completed"},"tags":[]}},{"cell_type":"code","source":"TRAIN_CONFIG = {\n    \"accelerator\": \"gpu\" if torch.cuda.is_available() else \"cpu\",\n    \"devices\": 1,\n    \"logger\": False,\n    \"enable_checkpointing\": torch.cuda.is_available(),\n    \"max_epochs\": 5,\n}","metadata":{"execution":{"iopub.status.busy":"2023-04-13T20:16:20.335944Z","iopub.execute_input":"2023-04-13T20:16:20.336799Z","iopub.status.idle":"2023-04-13T20:16:20.345305Z","shell.execute_reply.started":"2023-04-13T20:16:20.336771Z","shell.execute_reply":"2023-04-13T20:16:20.344134Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"dm = CustomDataModule()\nmodel = CycleGAN()\ntrainer = L.Trainer(**TRAIN_CONFIG)\n\ntrainer.fit(model, datamodule=dm)","metadata":{"papermill":{"duration":15422.135167,"end_time":"2023-03-19T06:14:51.028526","exception":false,"start_time":"2023-03-19T01:57:48.893359","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-04-13T20:16:20.346656Z","iopub.execute_input":"2023-04-13T20:16:20.347133Z","iopub.status.idle":"2023-04-13T21:24:01.109354Z","shell.execute_reply.started":"2023-04-13T20:16:20.347095Z","shell.execute_reply":"2023-04-13T21:24:01.108178Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Plotting the loss curves.","metadata":{"papermill":{"duration":0.145867,"end_time":"2023-03-19T06:14:51.360121","exception":false,"start_time":"2023-03-19T06:14:51.214254","status":"completed"},"tags":[]}},{"cell_type":"markdown","source":"To examine the loss curves of the generators and discriminators, we record the losses in `training_epoch_end` and define an additional method `loss_curves` above to plot the curves.","metadata":{"papermill":{"duration":0.15221,"end_time":"2023-03-19T06:14:51.681202","exception":false,"start_time":"2023-03-19T06:14:51.528992","status":"completed"},"tags":[]}},{"cell_type":"code","source":"model.loss_curves()","metadata":{"papermill":{"duration":0.651238,"end_time":"2023-03-19T06:14:52.570658","exception":false,"start_time":"2023-03-19T06:14:51.91942","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-04-13T21:24:01.112164Z","iopub.execute_input":"2023-04-13T21:24:01.112586Z","iopub.status.idle":"2023-04-13T21:24:01.530874Z","shell.execute_reply.started":"2023-04-13T21:24:01.112534Z","shell.execute_reply":"2023-04-13T21:24:01.529918Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"---","metadata":{"papermill":{"duration":0.147734,"end_time":"2023-03-19T06:14:52.876958","exception":false,"start_time":"2023-03-19T06:14:52.729224","status":"completed"},"tags":[]}},{"cell_type":"markdown","source":"# 4. Submission","metadata":{"papermill":{"duration":0.149119,"end_time":"2023-03-19T06:14:53.175117","exception":false,"start_time":"2023-03-19T06:14:53.025998","status":"completed"},"tags":[]}},{"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":{"papermill":{"duration":0.148192,"end_time":"2023-03-19T06:14:53.473205","exception":false,"start_time":"2023-03-19T06:14:53.325013","status":"completed"},"tags":[]}},{"cell_type":"code","source":"predictions = trainer.predict(model, datamodule=dm)","metadata":{"papermill":{"duration":126.035671,"end_time":"2023-03-19T06:16:59.658291","exception":false,"start_time":"2023-03-19T06:14:53.62262","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-04-13T21:24:01.532574Z","iopub.execute_input":"2023-04-13T21:24:01.533676Z","iopub.status.idle":"2023-04-13T21:24:03.851922Z","shell.execute_reply.started":"2023-04-13T21:24:01.533624Z","shell.execute_reply":"2023-04-13T21:24:03.850802Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Saving the generated images.","metadata":{"papermill":{"duration":0.146966,"end_time":"2023-03-19T06:16:59.953278","exception":false,"start_time":"2023-03-19T06:16:59.806312","status":"completed"},"tags":[]}},{"cell_type":"code","source":"os.makedirs(\"../images\", exist_ok=True)\n\nidx = 0\nfor tensor in predictions:\n    for monet in tensor:\n        save_image((monet.squeeze()*0.5+0.5), fp=f\"../images/{idx}.jpg\")\n        idx += 1\n\nshutil.make_archive(\"/kaggle/working/images\", \"zip\", \"/kaggle/images\")","metadata":{"papermill":{"duration":31.787796,"end_time":"2023-03-19T06:17:31.88707","exception":false,"start_time":"2023-03-19T06:17:00.099274","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-04-13T21:24:03.853774Z","iopub.execute_input":"2023-04-13T21:24:03.855308Z","iopub.status.idle":"2023-04-13T21:24:03.885766Z","shell.execute_reply.started":"2023-04-13T21:24:03.855261Z","shell.execute_reply":"2023-04-13T21:24:03.88466Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":" ---","metadata":{"papermill":{"duration":0.149773,"end_time":"2023-03-19T06:17:32.198249","exception":false,"start_time":"2023-03-19T06:17:32.048476","status":"completed"},"tags":[]}}]}