","metadata":{}},{"cell_type":"markdown","source":"### The Challenge: \nUnpaired Image-to-Image Translation using Cycle-Consistent Adversarial Networks\n\nA GAN consists of at least two neural networks: a generator model and a discriminator model. The generator is a neural network that creates the images. For our competition, you should generate images in the style of Monet. This generator is trained using a discriminator.\n\nThe two models will work against each other, with the generator trying to trick the discriminator, and the discriminator trying to accurately classify the real vs. generated images.\n\nYour task is to build a GAN that generates 7,000 to 10,000 Monet-style images.\n","metadata":{}},{"cell_type":"markdown","source":"#### Some references\n\n* [Kaggle Getting Started tutorial](https://www.kaggle.com/code/amyjang/monet-cyclegan-tutorial/notebook). The basic U-Net Generator\n* Original [paper](https://arxiv.org/abs/1406.2661) for Generative Adversarial Networks\n* Original [paper](https://arxiv.org/abs/1703.10593) and [code](https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix) for CycleGAN.To stabilize training, the author used an image buffer to store previously generated images to update the discriminator.\n* Original [paper](https://arxiv.org/abs/1611.04076) for LSGAN, which has shown to outperform BCE loss and is used in the original CycleGAN implementation. BCE loss may be susceptible to vanishing gradient problems and cause ineffective learning. For better training stability, we use LSGAN, which adopts the mean squared error for the adversarial criterion.\n* Pytorch implementations for CycleGAN [1](https://nn.labml.ai/gan/cycle_gan/index.html) and [2](https://github.com/eriklindernoren/PyTorch-GAN/blob/master/implementations/cyclegan/cyclegan.py).\n* [PyTorch documentation](https://pytorch.org/docs/stable/index.html)\n* [PyTorch Lightning documentation](https://lightning.ai/docs/pytorch/stable/starter/introduction.html)\n","metadata":{}},{"cell_type":"markdown","source":"---","metadata":{}},{"cell_type":"code","source":"# Fused Adam(FusedOptimizer)\n!pip install deepspeed","metadata":{"execution":{"iopub.status.busy":"2024-01-15T11:59:20.134032Z","iopub.execute_input":"2024-01-15T11:59:20.134842Z","iopub.status.idle":"2024-01-15T11:59:47.061709Z","shell.execute_reply.started":"2024-01-15T11:59:20.134806Z","shell.execute_reply":"2024-01-15T11:59:47.060741Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import glob\nimport os\nimport shutil\n\nimport deepspeed as ds\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport pytorch_lightning as L # lightweight PyTorch wrapper\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\n_ = L.seed_everything(0, workers=True) # sets a seed for reproducibility","metadata":{"execution":{"iopub.status.busy":"2024-01-15T12:00:55.891151Z","iopub.execute_input":"2024-01-15T12:00:55.891503Z","iopub.status.idle":"2024-01-15T12:01:03.865727Z","shell.execute_reply.started":"2024-01-15T12:00:55.891474Z","shell.execute_reply":"2024-01-15T12:01:03.864735Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 1. 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=(18, 8))\n plt.imshow(img_grid)\n plt.axis(\"off\")\n plt.title(title)\n plt.show()","metadata":{"execution":{"iopub.status.busy":"2024-01-15T12:01:06.441322Z","iopub.execute_input":"2024-01-15T12:01:06.442396Z","iopub.status.idle":"2024-01-15T12:01:06.448041Z","shell.execute_reply.started":"2024-01-15T12:01:06.442359Z","shell.execute_reply":"2024-01-15T12:01:06.446895Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Augmenting the images.","metadata":{}},{"cell_type":"markdown","source":"Before loading the datasets, we define a `CustomTransform` for image augmentation. This enhances the learning process by introducing more variaty into the images during training, instead of learning from the same set of images, especially when we only have 300 Monet paintings. The chosen image transformations include:\n* Scaling: Images are initially resized larger using the `Resize` operation and then randomly cropped back to the original size of 256 using `RandomCrop`. This process generates slightly different images to enrich the training dataset.\n* Horizontal Flipping: `RandomHorizontalFlip` is applied to randomly flip images horizontally. Since the content of both photos and Monet paintings is not significantly affected by horizontal orientation, this transformation adds variability to the dataset.\n* Color Variation: `ColorJitter` is used for randomly changing the colors of the images. This helps stimulate diverse lighting conditions for photos and introduces variability in the colors of Monet paintings.\n\nThese transformations are specifically applied during model training/fitting, as indicated by the `stage` argument. Finally, the images are scaled down for improved convergence during the training process.","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 \n # ensure images outside of training dataset are also of the same size\n self.transform = T.Resize((target_dim, target_dim), antialias=True)\n \n def __call__(self, img, stage):\n if stage == \"fit\":\n img = self.transform_train(img) # for training\n else:\n img = self.transform(img) # for validation/testing\n return img * 2 - 1 # normalizing the pixel values to be in [0,1]","metadata":{"execution":{"iopub.status.busy":"2024-01-15T12:01:08.764873Z","iopub.execute_input":"2024-01-15T12:01:08.765219Z","iopub.status.idle":"2024-01-15T12:01:08.772611Z","shell.execute_reply.started":"2024-01-15T12:01:08.76519Z","shell.execute_reply":"2024-01-15T12:01:08.771532Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Custom datasets.","metadata":{}},{"cell_type":"markdown","source":"This `CustomDataset` class is intended to be used for both the photos and Monet paintings datasets separately. Different instances of this class will be created for each dataset, and the datasets can later be combined while iterating through them during training or evaluation.","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 # normalizing pixel values\n return self.transform(img, stage=self.stage) # validation or testing stage","metadata":{"execution":{"iopub.status.busy":"2024-01-15T12:01:11.063129Z","iopub.execute_input":"2024-01-15T12:01:11.063727Z","iopub.status.idle":"2024-01-15T12:01:11.069348Z","shell.execute_reply.started":"2024-01-15T12:01:11.063696Z","shell.execute_reply":"2024-01-15T12:01:11.068462Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Iterating through the datasets.\n","metadata":{}},{"cell_type":"code","source":"DEBUG = False\n\nDM_CONFIG = { \n # the file paths\n \"monet_dir\": os.path.join(\"/kaggle/input/gan-getting-started/monet_jpg\", \"*.jpg\"),\n \"photo_dir\": os.path.join(\"/kaggle/input/gan-getting-started/photo_jpg\", \"*.jpg\"),\n \n \"loader_config\": {\n # the number of subprocesses \n \"num_workers\": os.cpu_count(),\n \n # enable faster data transfer to GPU during training\n \"pin_memory\": torch.cuda.is_available(),\n },\n \n # the validation/test batch size (for visualization purposes)\n \"sample_size\": 5,\n \n # the training/prediction batch size\n \"batch_size\": 1 if not DEBUG else 1,\n}","metadata":{"execution":{"iopub.status.busy":"2024-01-15T12:01:13.209052Z","iopub.execute_input":"2024-01-15T12:01:13.209411Z","iopub.status.idle":"2024-01-15T12:01:13.215316Z","shell.execute_reply.started":"2024-01-15T12:01:13.209376Z","shell.execute_reply":"2024-01-15T12:01:13.214243Z"},"trusted":true},"outputs":[],"execution_count":null},{"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.\n\nThe validation, test and prediction datasets will only contain photos for simplicity since we are only trying to generate Monet-style images.\n\nTo organize the data-related steps needed for a PyTorch Lightning model, such as loading datasets and setting up data loaders we define a custom`LightningDataModule`. The key components:\n* `setup` to create the different instances of `CustomDataset` defined above.\n* `train_dataloader` to load both the Monet paintings and photos for training.\n* `val_dataloader` to load only 1 batch of photos and use it repeatedly to visualize the progress of adding Monet style at different points of training.\n* `test_dataloader` to load a small subset of photos to test the results of photo-to-Monet translation on different sample photos after training. \n* `predict_dataloader` to load the entire dataset of photos and add Monet style for submission.\n","metadata":{}},{"cell_type":"code","source":"class 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 \n # store file paths\n self.monet_filenames = sorted(glob.glob(monet_dir))\n self.photo_filenames = sorted(glob.glob(photo_dir))\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":"2024-01-15T12:01:24.15214Z","iopub.execute_input":"2024-01-15T12:01:24.152495Z","iopub.status.idle":"2024-01-15T12:01:24.16351Z","shell.execute_reply.started":"2024-01-15T12:01:24.152465Z","shell.execute_reply":"2024-01-15T12:01:24.162616Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Visual inspection and validation of the data augmentation process. It helps to ensure that the custom data module, datasets, and data loaders are correctly set up and that the image augmentation transformations are being applied as expected.","metadata":{}},{"cell_type":"code","source":"# set batch_size to 5\ndm_sample = CustomDataModule(batch_size=5, **{k: v for k, v in DM_CONFIG.items() if k != \"batch_size\"})\ndm_sample.setup(\"fit\")\n\ntrain_loader = dm_sample.train_dataloader()\nimgs = next(iter(train_loader))\n","metadata":{"execution":{"iopub.status.busy":"2024-01-15T12:01:26.80601Z","iopub.execute_input":"2024-01-15T12:01:26.806688Z","iopub.status.idle":"2024-01-15T12:01:28.040243Z","shell.execute_reply.started":"2024-01-15T12:01:26.806639Z","shell.execute_reply":"2024-01-15T12:01:28.037588Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 2. Building CycleGAN Architecture","metadata":{}},{"cell_type":"markdown","source":"## Generator arhitectures:","metadata":{}},{"cell_type":"markdown","source":"### U-Net \nA 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.\n\n*Example of the U-Net architecture [[source](https://paperswithcode.com/method/u-net)].*","metadata":{}},{"cell_type":"markdown","source":"### ResNet\nIn the original CycleGAN implementation, the author used a ResNet generator. 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.\n\n*Example of the ResNet architecture [[source](https://paperswithcode.com/method/resnet)].*","metadata":{}},{"cell_type":"markdown","source":"### Downsampling blocks\nDownsampling blocks are used to reduce the spatial dimensions of an input image while increasing the number of feature maps. This helps in capturing hierarchical features in the image.","metadata":{}},{"cell_type":"code","source":"class Downsampling(nn.Module):\n def __init__(\n self,\n in_channels,\n out_channels,\n kernel_size=4,\n stride=2,\n padding=1,\n norm=True,\n lrelu=True,\n ):\n super().__init__()\n self.block = nn.Sequential(\n nn.Conv2d(in_channels, out_channels,\n kernel_size=kernel_size, stride=stride, padding=padding, bias=not norm),\n )\n if norm:\n self.block.append(nn.InstanceNorm2d(out_channels, affine=True))\n if lrelu is not None:\n self.block.append(nn.LeakyReLU(0.2, True) if lrelu else nn.ReLU(True))\n \n def forward(self, x):\n return self.block(x)","metadata":{"execution":{"iopub.status.busy":"2024-01-15T12:01:30.494667Z","iopub.execute_input":"2024-01-15T12:01:30.495559Z","iopub.status.idle":"2024-01-15T12:01:30.506046Z","shell.execute_reply.started":"2024-01-15T12:01:30.495504Z","shell.execute_reply":"2024-01-15T12:01:30.505016Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Upsampling blocks\nUpsampling blocks are used to increase the spatial dimensions of an input tensor while reducing the number of feature maps, eventually reconstructing an image with the original size.","metadata":{}},{"cell_type":"code","source":"class Upsampling(nn.Module):\n def __init__(\n self,\n in_channels,\n out_channels,\n kernel_size=4,\n stride=2,\n padding=1,\n output_padding=0,\n dropout=False,\n ):\n super().__init__()\n self.block = nn.Sequential(\n nn.ConvTranspose2d(in_channels, out_channels,\n 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(True))\n \n def forward(self, x):\n return self.block(x)","metadata":{"execution":{"iopub.status.busy":"2024-01-15T12:01:36.194411Z","iopub.execute_input":"2024-01-15T12:01:36.194768Z","iopub.status.idle":"2024-01-15T12:01:36.202159Z","shell.execute_reply.started":"2024-01-15T12:01:36.19474Z","shell.execute_reply":"2024-01-15T12:01:36.201222Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Residual Block\nBy adding the original input to the output, the residual block allows the network to learn to modify the original input rather than reconstructing the entire output from scratch. This helps with the flow of gradients during backpropagation, enabling the training of deeper networks.","metadata":{}},{"cell_type":"code","source":"class ResBlock(nn.Module):\n def __init__(self, in_channels, kernel_size=3, padding=1):\n super().__init__()\n self.block = nn.Sequential(\n nn.ReflectionPad2d(padding),\n Downsampling(in_channels, in_channels,\n kernel_size=kernel_size, stride=1, padding=0, lrelu=False),\n nn.ReflectionPad2d(padding),\n Downsampling(in_channels, in_channels,\n kernel_size=kernel_size, stride=1, padding=0, lrelu=None),\n )\n \n def forward(self, x):\n return x + self.block(x)","metadata":{"execution":{"iopub.status.busy":"2024-01-15T12:01:38.841142Z","iopub.execute_input":"2024-01-15T12:01:38.841509Z","iopub.status.idle":"2024-01-15T12:01:38.848342Z","shell.execute_reply.started":"2024-01-15T12:01:38.841479Z","shell.execute_reply":"2024-01-15T12:01:38.847397Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class UNetGenerator(nn.Module):\n def __init__(self, hid_channels, in_channels, out_channels):\n super().__init__()\n self.downsampling_path = nn.Sequential(\n Downsampling(in_channels, hid_channels, norm=False), # 64x128x128 out_channels-height-width\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, instance norm does not work on 1x1\n )\n self.upsampling_path = 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(), # hyperbolic tangent\n )\n \n def forward(self, x):\n # Downsampling through the model\n skips = []\n for down in self.downsampling_path:\n x = down(x)\n skips.append(x)\n skips = reversed(skips[:-1])\n \n # Upsampling and establishing the skip connections\n for up, skip in zip(self.upsampling_path, skips):\n x = up(x)\n x = torch.cat([x, skip], dim=1)\n return self.feature_block(x)\n \nclass ResNetGenerator(nn.Module):\n def __init__(self, hid_channels, in_channels, out_channels, num_resblocks):\n super().__init__()\n self.model = nn.Sequential(\n # downsampling path\n nn.ReflectionPad2d(3), # to handle border pixels\n Downsampling(in_channels, hid_channels,\n kernel_size=7, stride=1, padding=0, lrelu=False), # 64x256x256\n Downsampling(hid_channels, hid_channels*2, kernel_size=3, lrelu=False), # 128x128x128\n Downsampling(hid_channels*2, hid_channels*4, kernel_size=3, lrelu=False), # 256x64x64\n \n # residual blocks\n *[ResBlock(hid_channels*4) for _ in range(num_resblocks)], # 256x64x64\n \n # upsampling path\n Upsampling(hid_channels*4, hid_channels*2, kernel_size=3, output_padding=1), # 128x128x128\n Upsampling(hid_channels*2, hid_channels, kernel_size=3, output_padding=1), # 64x256x256\n nn.ReflectionPad2d(3), # to handle border pixels\n nn.Conv2d(hid_channels, out_channels, kernel_size=7, stride=1, padding=0), # 3x256x256\n nn.Tanh(),# pixels in the range [-1,1]\n )\n \n def forward(self, x):\n return self.model(x)\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(hid_channels, in_channels, out_channels, num_resblocks)\n else:\n raise NotImplementedError(f\"Generator name '{gen_name}' not recognized.\")","metadata":{"execution":{"iopub.status.busy":"2024-01-15T12:01:41.572296Z","iopub.execute_input":"2024-01-15T12:01:41.572997Z","iopub.status.idle":"2024-01-15T12:01:41.591224Z","shell.execute_reply.started":"2024-01-15T12:01:41.572966Z","shell.execute_reply":"2024-01-15T12:01:41.590283Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## PatchGAN discriminator\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.\n\nThe idea behind using a PatchGAN discriminator is to provide fine-grained feedback to the generator by evaluating the realism of local patches in the generated image. This allows the generator to focus on capturing and improving details at a local level rather than the entire image.\n\n*PatchGAN discriminator [[source]](https://paperswithcode.com/method/patchgan).*","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 (num_channels-h-w)\n )# 1 channel for binary classification task, 30-30 spatial dimensions of the feature map\n \n def forward(self, x):\n return self.block(x)","metadata":{"execution":{"iopub.status.busy":"2024-01-15T12:01:45.672728Z","iopub.execute_input":"2024-01-15T12:01:45.673106Z","iopub.status.idle":"2024-01-15T12:01:45.679768Z","shell.execute_reply.started":"2024-01-15T12:01:45.673077Z","shell.execute_reply":"2024-01-15T12:01:45.678745Z"},"trusted":true},"outputs":[],"execution_count":null},{"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) # to add batch dimension\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":"2024-01-15T12:01:48.67789Z","iopub.execute_input":"2024-01-15T12:01:48.678255Z","iopub.status.idle":"2024-01-15T12:01:48.687224Z","shell.execute_reply.started":"2024-01-15T12:01:48.678227Z","shell.execute_reply":"2024-01-15T12:01:48.686266Z"},"trusted":true},"outputs":[],"execution_count":null},{"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\": 9,\n \n # the number of filters in the first layer for the generators and discriminators\n \"hid_channels\": 64,\n \n # using DeepSpeed's FusedAdam \n \"optimizer\": ds.ops.adam.FusedAdam if torch.cuda.is_available() else torch.optim.Adam,\n \n # the learning rate and beta parameters for the Adam optimizer\n \"lr\": 2e-4,\n \"betas\": (0.5, 0.999),\n \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 \n # the size of the buffer that stores previously generated images\n \"buffer_size\": 100,\n \n # the number of epochs for training\n \"num_epochs\": 19 if not DEBUG else 2,\n \n # the number of epochs before starting the learning rate decay\n \"decay_epochs\": 19 if not DEBUG else 1,\n}","metadata":{"execution":{"iopub.status.busy":"2024-01-15T12:18:30.710747Z","iopub.execute_input":"2024-01-15T12:18:30.711156Z","iopub.status.idle":"2024-01-15T12:18:30.718321Z","shell.execute_reply.started":"2024-01-15T12:18:30.711124Z","shell.execute_reply":"2024-01-15T12:18:30.717294Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### CycleGAN\n\nIn a CycleGAN, two generators (`gen_PM` and `gen_MP`) and two discriminators (`disc_M` and `disc_P`) are employed to facilitate photo-to-Monet and Monet-to-photo translations. These generators and discriminators are optimized using the Adam optimizer during training. To guide the training process, specific loss functions are defined:\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.\n\nFrom the above, the mean squared error and the l1 loss are defined as the **adversarial criterion** and the **reconstruction criterion** respectively.","metadata":{}},{"cell_type":"markdown","source":"### Building the CycleGAN model.","metadata":{}},{"cell_type":"markdown","source":"We structure the modeling code by defining loss functions and incorporating essential methods within `LightningModule`. These methods include `forward` for generating Monet-style images, `setup` for initializing weights and biases, `configure_optimizers` for setting up optimizers and learning rate schedules, `training_step` for computing generator and discriminator losses, `validation_step` for visualizing progress during training, `test_step` for evaluating results on sample photos, `predict_step` for running the forward method during prediction, and various methods for logging and reporting during training and prediction epochs.","metadata":{}},{"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 isinstance(m, (nn.Conv2d, nn.ConvTranspose2d, nn.InstanceNorm2d)):\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":"2024-01-15T12:18:33.484305Z","iopub.execute_input":"2024-01-15T12:18:33.484696Z","iopub.status.idle":"2024-01-15T12:18:33.520568Z","shell.execute_reply.started":"2024-01-15T12:18:33.484645Z","shell.execute_reply":"2024-01-15T12:18:33.519637Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 3. Model Training","metadata":{}},{"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\": 6 if not DEBUG else 1,\n}","metadata":{"execution":{"iopub.status.busy":"2024-01-15T12:18:37.895708Z","iopub.execute_input":"2024-01-15T12:18:37.896604Z","iopub.status.idle":"2024-01-15T12:18:37.903108Z","shell.execute_reply.started":"2024-01-15T12:18:37.896568Z","shell.execute_reply":"2024-01-15T12:18:37.902158Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"`Trainer` automatically handle the training loop by running the fit method.","metadata":{}},{"cell_type":"code","source":"model = CycleGAN(**MODEL_CONFIG)","metadata":{"execution":{"iopub.status.busy":"2024-01-15T12:18:40.788175Z","iopub.execute_input":"2024-01-15T12:18:40.788843Z","iopub.status.idle":"2024-01-15T12:18:41.704472Z","shell.execute_reply.started":"2024-01-15T12:18:40.788808Z","shell.execute_reply":"2024-01-15T12:18:41.703453Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"dm = CustomDataModule(**DM_CONFIG)\nmodel = CycleGAN(**MODEL_CONFIG)\ntrainer = L.Trainer(**TRAIN_CONFIG)\n\ntrainer.fit(model, datamodule=dm)","metadata":{"execution":{"iopub.status.busy":"2024-01-14T13:14:10.40495Z","iopub.execute_input":"2024-01-14T13:14:10.405366Z","iopub.status.idle":"2024-01-14T13:56:55.17526Z","shell.execute_reply.started":"2024-01-14T13:14:10.405335Z","shell.execute_reply":"2024-01-14T13:56:55.174288Z"},"trusted":true},"outputs":[],"execution_count":null},{"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":"2024-01-14T13:58:23.26034Z","iopub.execute_input":"2024-01-14T13:58:23.261374Z","iopub.status.idle":"2024-01-14T13:58:29.964937Z","shell.execute_reply.started":"2024-01-14T13:58:23.261332Z","shell.execute_reply":"2024-01-14T13:58:29.963768Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 4. Submission","metadata":{}},{"cell_type":"markdown","source":"`Trainer` automatically computes the predictions by running the `predict` method.","metadata":{}},{"cell_type":"code","source":"predictions = trainer.predict(model, datamodule=dm)","metadata":{"execution":{"iopub.status.busy":"2024-01-14T13:58:44.958878Z","iopub.execute_input":"2024-01-14T13:58:44.959295Z","iopub.status.idle":"2024-01-14T13:59:40.961318Z","shell.execute_reply.started":"2024-01-14T13:58:44.959262Z","shell.execute_reply":"2024-01-14T13:59:40.960328Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Saving the generated images.","metadata":{}},{"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(\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":{"execution":{"iopub.status.busy":"2024-01-14T13:59:46.028995Z","iopub.execute_input":"2024-01-14T13:59:46.029429Z","iopub.status.idle":"2024-01-14T14:00:08.065905Z","shell.execute_reply.started":"2024-01-14T13:59:46.029397Z","shell.execute_reply":"2024-01-14T14:00:08.064895Z"},"trusted":true},"outputs":[],"execution_count":null}]}
