{"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":"# University of Washington CSE 493G1 Final Project: MonetGAN - Generative Art\n## Garrett Devereux, Deekshita S Doli, and Begoña García Malaxechebarría\n\nThis notebook serves as a baseline for our Generative Art project utilizing the [Kaggle I'm Something of a Painter Myself](https://www.kaggle.com/competitions/gan-getting-started/data) dataset. Here, we attempt to create a simple CycleGAN and seek to achieve non-trivial results. Through this process, we will become more comfortable coding, training, and tuning our GAN with initial results we can work to improve in the coming weeks.\n\n\nThis notebook is heavily inspired from a variety of different references:\n\n*   The Kaggle suggested [CycleGAN Tutorial Notebook](https://www.kaggle.com/code/amyjang/monet-cyclegan-tutorial/notebook)\n*   [PyTorch Lightning Notebook](https://www.kaggle.com/code/chongzhenjie/photo-to-monet-cyclegan-pytorch-lightning/notebook) that introduces the implementation and benefits of training with PyTorch Lightning. \n* [Original CycleGAN paper](https://arxiv.org/abs/1703.10593)\n* [Original UNet Generator paper](https://paperswithcode.com/paper/u-net-convolutional-networks-for-biomedical)\n* [Original PatchGAN Discriminator paper](https://paperswithcode.com/method/patchgan)\n","metadata":{"id":"X7IHgrsqd-W0"}},{"cell_type":"markdown","source":"## Import Necessary Libraries","metadata":{"id":"YYSTgmPRp8xm"}},{"cell_type":"markdown","source":"Ensure PyTorch Lightning is installed along with all necessary libraries","metadata":{"id":"U4i8V1W1FHC7"}},{"cell_type":"code","source":"!pip install pytorch_lightning==1.9.3;","metadata":{"executionInfo":{"elapsed":19746,"status":"ok","timestamp":1684109177519,"user":{"displayName":"Garrett Devereux","userId":"06903598260790252274"},"user_tz":420},"id":"-mQnQb6rm1C7","outputId":"fa2b5398-9206-4c8d-aee3-caeeec4ab8c3","trusted":true},"execution_count":null,"outputs":[]},{"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 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\ndevice = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\nprint(device)","metadata":{"executionInfo":{"elapsed":5603,"status":"ok","timestamp":1684109217082,"user":{"displayName":"Garrett Devereux","userId":"06903598260790252274"},"user_tz":420},"id":"zXMDumZap7i4","outputId":"f90b07e9-b9b1-4742-bd0e-fab18bdd746f","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Load and Preprocess the Data\n\nAs this is the baseline example, we won't worry about Data Augmentation or any advanced preprocessing techniques. Instead we will just focus on setting up the Dataset and Dataloaders as seen in [PyTorch Tutorial from Recitation](https://courses.cs.washington.edu/courses/cse493g1/23sp/schedule/lectures/recitation4.pdf) and try to Visualize the Dataset ","metadata":{"id":"5_LMm626Iz3r"}},{"cell_type":"markdown","source":"### Storing The Dataset\n\nIn order to load and store the Dataset, we need to define a ```CustomImageDataset``` class that implements three functions:\n\n\n*   **```__init__```**: store any internal values that may be of use\n*   **```__len__```**: get the size of the dataset\n*   **```__getitem__```**: get the i'th image in the dataset\n\nDifferent instances of this class will be used to store the Photos and the Monet Paintings separately\n\n\n\n","metadata":{"id":"NZOk1fMqKGz3"}},{"cell_type":"code","source":"class CustomImageDataset(Dataset):\n    def __init__(self, filenames, stage):\n        # A list of every filename of every image in the dataset\n        self.filenames = filenames\n        # Stage of the current dataset (fit, test, predict)\n        # Helps determine correct behavior when using the Dataset\n        self.stage = stage\n        \n    def __len__(self):\n        # Return the length of the list\n        return len(self.filenames)\n    \n    def __getitem__(self, idx):\n        # Get the corresponding image and scale it\n        img_name = self.filenames[idx]\n        img = read_image(img_name) / 255.0\n        return img * 2 - 1","metadata":{"id":"eZGw_PPTKLh2","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### DataLoaders and DataModules\n\nNext, we prepare the Dataset by separately loading the Monet Paintings and Photos into their own DataLoaders. As we are going to use PyTorch Lightning for the training process, we will handle all of this in a single Custom ```LightningDataModule```. This allows us to greatly simplify the configuration while encapsulating the five steps of preprocessing:\n\n\n\n\n1.   Download / tokenize / process.\n2.   Clean and (maybe) save to disk.\n3.   Load inside ```Dataset```.\n4.   Apply transforms (rotate, tokenize, etc…).\n5.   Wrap inside a ```DataLoader```.\n\n\n\n\n\n\nSee the [PyTorch Lightning Documentation](https://lightning.ai/docs/pytorch/stable/data/datamodule.html) for more details.","metadata":{"id":"YNnbFRg1MdTD"}},{"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        # Store Configuration details to pass to DataLoaders\n        self.loader_config = loader_config\n        # Batch Size for Val and Test stages (mainly for visualization)\n        self.sample_size = sample_size\n        # Batch Size for Predict phase\n        self.batch_size = batch_size       \n            \n        # Store Sorted filenames of all the Monet Paintings and Photos in the dataset\n        self.monet_filenames = sorted(glob.glob(monet_dir))\n        self.photo_filenames = sorted(glob.glob(photo_dir))\n        \n    # Handle the Creation of Datasets and Train/Test/Val split\n    def setup(self, stage):\n        if stage == \"fit\":\n            self.train_monet = CustomImageDataset(self.monet_filenames, stage)\n            self.train_photo = CustomImageDataset(self.photo_filenames, stage)\n            \n        if stage in [\"fit\", \"test\", \"predict\"]:\n            # Validation is the dataset of Photos that will be used\n            # to be converted to Monet Paintings\n            self.valid_photo = CustomImageDataset(self.photo_filenames, None)\n\n    # After setup is called, train_dataloader will return the Wrapped Loaders\n    # Of the Monet Paintings AND Photos for training\n    def train_dataloader(self):\n        # Create config for loaders with passed in config\n        loader_config = {\n            \"shuffle\": True,\n            \"drop_last\": True,\n            \"batch_size\": self.batch_size,\n            **self.loader_config,\n        }\n        # Create Dataloaders for each Dataset and return\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 loaders\n    \n    # Test and Val Dataloaders are the same\n    # Just the Dataset of Photos\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    # Predict also uses Photos dataset \n    # but uses batch size rather than sample size\n    def predict_dataloader(self):\n        return DataLoader(\n            self.valid_photo,\n            batch_size=self.batch_size,\n            **self.loader_config,\n        )","metadata":{"id":"Yno-7MPvMOhI","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Configuration Details\n\nHere we setup the configuration that allows us to create our ```CustomDataModule``` setting key values such as the directories to the datasets, the loader config, and the batch sizes.","metadata":{"id":"2m5f0k1BUkiM"}},{"cell_type":"code","source":"DM_CONFIG = {    \n    # Directories of the Monet Paintings and Photos\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 (excluding the main process) used for data loading\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 (mainly for visualization purposes)\n    \"sample_size\": 5,\n    \n    # the training/prediction batch size\n    \"batch_size\": 1\n}","metadata":{"id":"Zll8B3WxQv0U","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Understanding the Dataset\n\nNow we instantiate the ```CustomDataModule``` using the above configuration, then call ```train_dataloader()``` to get the Monet and Photos Dataloaders. Looking at the data that is produced, we can see that with a batch size of 1 we have 300 Monet Paintings and 7028 Photos. Additionally, we can see that our images are shape 3 x 256 x 256 where we each photo has 3 channels with height and width of 256","metadata":{"id":"9Aufcf2z0aDE"}},{"cell_type":"code","source":"# Instiantiate the CustomDataModule with the DM Config\ndm_understand = CustomDataModule(**{k: v for k, v in DM_CONFIG.items()})\n# Call setup with stage fit to setup both the Monet Loader and the Photos Loader\ndm_understand.setup(\"fit\")\n\n# Train_loader now returns both the Monet Loader and the Photos Loader\ntrain_loader = dm_understand.train_dataloader()\nmonet_ds = train_loader[\"monet\"]\nphoto_ds = train_loader[\"photo\"]\n\n# Lengths of the Dataset\nprint(\"Number of Monet Paintings: \" + str(len(monet_ds)))\nprint(\"Number of Photos: \" + str(len(photo_ds)))\nprint()\n\n# Image size 3 x 256 x 256\nprint(\"Image Sizes: \")\nprint(next(iter(monet_ds))[0].shape)\nprint(next(iter(photo_ds))[0].shape)\nprint()\n\n# Store a set of example images to understand and visualize\nexample_monets = torch.empty((5, 3, 256, 256))\nexample_photos = torch.empty((5, 3, 256, 256))\n\n# Example of how to keep getting images form each Dataset\n# Note Batch Size = 1 so we get a batch of 1 on each next()\nfor i in range(5):\n  example_monets[i] = next(iter(monet_ds))\n  example_photos[i] = next(iter(photo_ds))\n\nprint(\"Example Batch Sizes:\")\nprint(example_monets.shape)\nprint(example_photos.shape)","metadata":{"executionInfo":{"elapsed":51295,"status":"ok","timestamp":1684109334569,"user":{"displayName":"Garrett Devereux","userId":"06903598260790252274"},"user_tz":420},"id":"GnmReHHjQsZF","outputId":"16f0c5f9-9289-4831-f0e9-419fd7152ce1","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Visualizing the Dataset\n\nWith an understanding of how we can initialize the Dataloaders, iterate through batches, and visualize the shapes of the images and batches, let's now look at what the images actually look like:","metadata":{"id":"woPcoI_c08yw"}},{"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()\n\nshow_img(example_monets, nrow=5, title=\"Monet Paintings\")\nshow_img(example_photos, nrow=5, title=\"Photos\")","metadata":{"executionInfo":{"elapsed":12201,"status":"ok","timestamp":1684109375499,"user":{"displayName":"Garrett Devereux","userId":"06903598260790252274"},"user_tz":420},"id":"SecQv3LeRnb-","outputId":"962e2f64-a8c6-47da-86f2-8690ef94e638","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Generators and Discriminators in GANs\n\n![image.jpg](https://cpang4.github.io/gan/content/gen-dis.jpg)\n[Image Source](https://cpang4.github.io/gan/)\n\n\n[From: A Gentle Introduction to CycleGAN for Image Translation](https://machinelearningmastery.com/what-is-cyclegan/):\n\nThe GAN architecture is an approach to training a model for image synthesis that is comprised of two models: a generator model and a discriminator model. The generator takes a point from a latent space as input and generates new plausible images from the domain, and the discriminator takes an image as input and predicts whether it is real (from a dataset) or fake (generated). Both models are trained in a game, such that the generator is updated to better fool the discriminator and the discriminator is updated to better detect generated images.\n\nThe CycleGAN is an extension of the GAN architecture that involves the simultaneous training of two generator models and two discriminator models.\n\nOne generator takes images from the first domain as input and outputs images for the second domain, and the other generator takes images from the second domain as input and generates images for the first domain. Discriminator models are then used to determine how plausible the generated images are and update the generator models accordingly.\n","metadata":{"id":"HGZqaIka7sty"}},{"cell_type":"markdown","source":"## Generator\n\nFrom the [Tutorial Notebook](https://www.kaggle.com/code/amyjang/monet-cyclegan-tutorial/notebook):\n\n\"We'll be using a UNET architecture for our generators in our CycleGAN. To build our generator, let's first define our downsample and upsample methods.\n\nThe downsample, as the name suggests, reduces the 2D dimensions, the width and height, of the image by the stride. The stride is the length of the step the filter takes. Since the stride is 2, the filter is applied to every other pixel, hence reducing the weight and height by 2.\n\nWe'll be using an instance normalization instead of batch normalization.\"\n\nUNET - Architecture\n\n![image.png](https://lmb.informatik.uni-freiburg.de/people/ronneber/u-net/u-net-architecture.png)\n[Image Source](https://paperswithcode.com/method/u-net)","metadata":{"id":"y9SdUuDvAQ5P"}},{"cell_type":"markdown","source":"For simplicity, we will initialize all the weights in the Generator with a Random Normal. More advanced initialization techniques can be looked into later (i.e. Xavier or Kaiming init)","metadata":{"id":"QwreT8nAbVRr"}},{"cell_type":"code","source":"def init_weights(m):\n    if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d, nn.InstanceNorm2d)):\n      torch.nn.init.normal_(m.weight, mean=0.0, std=0.02)\n      if m.bias is not None:\n        nn.init.constant_(m.bias, 0.0)","metadata":{"id":"R87lWdITbgVb","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#### Downsampling Module\nBy initializing model with in ```in_channels``` and ```out_channels``` then calling ```forward()```, we can take an input of size ```(batch_size, in_channels, H, W)``` and output a size of ```(batch_size, out_channels, H/2, W/2)```.","metadata":{"id":"7QBaIV6yCx6J"}},{"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        # Apply Init Weights\n        self.block.apply(init_weights)\n        \n    def forward(self, x):\n        return self.block(x)","metadata":{"id":"45T4Jv_1ckmk","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Quick Sanity Check to Ensure Downsampling is working correctly. Note we can also enable Instance Normalization by adding ```norm=True``` in the initialization of our Downsampling model.","metadata":{"id":"JtbD1FXPC8Ya"}},{"cell_type":"code","source":"def test_Downsampling():\n  # Test going from 3 in_channels -> 64 hidden_channels -> 128 out channels\n  # with batch size of 5 and image size 3 x 256 x 256\n  in_channels = 3\n  hidden_channels = 64\n  out_channels = 128\n  batch_size = 5\n  H = 256\n  W = 256\n  # Create Zero Tensor of desired input shape\n  x = torch.zeros(((batch_size, in_channels, H, W)))\n  print(\"Initial Size:\")\n  print(x.shape)\n  print()\n\n  # Downsample 1: define beginning in->out channels\n  model1 = Downsampling(in_channels, hidden_channels)\n  scores = model1(x)\n  # You should see size [batch_size, hidden_channels, H/2, W/2]\n  # Which = [5, 64, 128, 128] in this case\n  print(\"Downsample 1:\")\n  print(scores.size())\n  print()\n\n  # Downsampling 2, need to redefine input/output size\n  model2 = Downsampling(hidden_channels, out_channels)\n  scores2 = model2(scores)\n  # You should see size [batch_size, out_channels, H/4, W/4]\n  # Which = [5, 128, 64, 64] in this case\n  print(\"Downsample 2:\")\n  print(scores2.size())\n\n\ntest_Downsampling()","metadata":{"executionInfo":{"elapsed":904,"status":"ok","timestamp":1684109404290,"user":{"displayName":"Garrett Devereux","userId":"06903598260790252274"},"user_tz":420},"id":"3AGFG97cB7yh","outputId":"423a72d7-b78b-48cc-8868-f28ee796eedf","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#### Upsampling Module\nBy initializing model with in ```in_channels``` and ```out_channels``` then calling ```forward()```, we can take an input of size ```(batch_size, in_channels, H, W)``` and output a size of ```(batch_size, out_channels, H*2, W*2)```","metadata":{"id":"xVnpgUcWGdV8"}},{"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        # Apply Init Weights\n        self.block.apply(init_weights)\n        \n    def forward(self, x):\n        return self.block(x)","metadata":{"id":"a_KIiJc4cu3d","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Quick Sanity Check to Ensure Upsampling is working correctly. Note we can also enable dropout by adding ```dropout=True``` in the initialization of our Upsampling model.","metadata":{"id":"yIEjVxb4HviY"}},{"cell_type":"code","source":"def test_Upsampling():\n  # Test going from 128 in_channels -> 64 hidden_channels -> 3 out channels\n  # with batch size of 5 and input size 128 x 64 x 64\n  in_channels = 128\n  hidden_channels = 64\n  out_channels = 3\n  batch_size = 5\n  H = 64\n  W = 64\n  # Create Zero Tensor of desired input shape\n  x = torch.zeros(((batch_size, in_channels, H, W)))\n  print(\"Initial Size:\")\n  print(x.shape)\n  print()\n\n  # Upsample 1: define beginning in->out channels\n  model1 = Upsampling(in_channels, hidden_channels)\n  scores = model1(x)\n  # You should see size [batch_size, hidden_channels, H*2, W*2]\n  # Which = [5, 64, 128, 128] in this case\n  print(\"Upsample 1:\")\n  print(scores.size())\n  print()\n\n  # Upsampling 2, need to redefine input/output size\n  model2 = Upsampling(hidden_channels, out_channels)\n  scores2 = model2(scores)\n  # You should see size [batch_size, out_channels, H*4, W*4]\n  # Which = [5, 3, 256, 256] in this case\n  print(\"Upsample 2:\")\n  print(scores2.size())\n\n\ntest_Upsampling()","metadata":{"executionInfo":{"elapsed":1226,"status":"ok","timestamp":1684109413177,"user":{"displayName":"Garrett Devereux","userId":"06903598260790252274"},"user_tz":420},"id":"SECTyFKqHvG4","outputId":"b1474008-0246-4d02-a0a6-c5dc515bdb2e","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Build the Generator!\n\nThe generator first downsamples the input image and then upsample while establishing long skip connections, giving it the U-shaped architecture. Skip connections are a way to help bypass the vanishing gradient problem by concatenating the output of a layer to multiple layers instead of only one. Here we concatenate the output of the downsample layer to the upsample layer in a symmetrical fashion. \n\nNote that we initialize the UNet Generator weights with Random Normals for simplicity, more advanced init techniques may be looked into layer (i.e. Xavier or Kaiming init...)","metadata":{"id":"sUsii0Ihcz-J"}},{"cell_type":"code","source":"class UNetGenerator(nn.Module):\n    def __init__(self, hid_channels, in_channels, out_channels):\n        super().__init__()\n\n        # Initialize Series Downsampling blocks mapping from an initial\n        # (batch_size, 3, 256, 256) -> (batch_size, 512, 1, 1)\n        self.downsampling_path = nn.Sequential(\n            # Note all following have batch_size as 0 dimension\n            # And are assuming hid=64, in=3, out=3, H=W=256\n            # However this function will generalize for other sizes\n            # (3, 256, 256) -> (64, 128, 128)\n            Downsampling(in_channels, hid_channels, norm=False),\n            # (64, 128, 128) -> (128, 64, 64)\n            Downsampling(hid_channels, hid_channels*2),\n            # (128, 64, 64) -> (256, 32, 32)\n            Downsampling(hid_channels*2, hid_channels*4),\n            # (256, 32, 32) -> (512, 16, 16)\n            Downsampling(hid_channels*4, hid_channels*8),\n            # (512, 16, 16) -> (512, 8, 8)\n            Downsampling(hid_channels*8, hid_channels*8),\n            # (512, 8, 8) -> (512, 4, 4)\n            Downsampling(hid_channels*8, hid_channels*8),\n            # (512, 4, 4) -> (512, 2, 2)\n            Downsampling(hid_channels*8, hid_channels*8),\n            # (512, 2, 2) -> (512, 1, 1) (turn off instance norm for size 1)\n            Downsampling(hid_channels*8, hid_channels*8, norm=False),\n        )\n\n        # Intialize Series of Upsamping blocks mapping from\n        # (batch_size, 512, 1, 1) -> (batch_size, 128, 128, 128)\n        self.upsampling_path = nn.Sequential(\n            # Note all following have batch_size as 0 dimension\n            # And are assuming hid=64, with initial input (512, 1, 1)\n            # However this function will generalize for other sizes\n            # (512, 1, 1) -> (512, 2, 2)\n            Upsampling(hid_channels*8, hid_channels*8, dropout=True),\n            # Concat with Down => (1024, 2, 2) -> (512, 4, 4)\n            Upsampling(hid_channels*16, hid_channels*8, dropout=True),\n            # Concat with Down => (1024, 4, 4) -> (512, 8, 8)\n            Upsampling(hid_channels*16, hid_channels*8, dropout=True),\n            # Concat with Down => (1024, 8, 8) -> (512, 16, 16)\n            Upsampling(hid_channels*16, hid_channels*8),\n            # Concat with Down => (1024, 16, 16) -> (256, 32, 32)\n            Upsampling(hid_channels*16, hid_channels*4), # (256+256)x32x32\n            # Concat with Down => (512, 32, 32) -> (128, 64, 64)\n            Upsampling(hid_channels*8, hid_channels*2),\n            # Concat with Down => (256, 64, 64) -> (64, 128, 128)\n            Upsampling(hid_channels*4, hid_channels),\n            # Concat with Down for final shape (128, 128, 128)\n        )\n\n        # Final Feature Block with nonlinearity to map back to image shape\n        # Input Size (batch_size, 128, 128, 128) -> (batch_size, 3, 256, 256)\n        self.feature_block = nn.Sequential(\n            nn.ConvTranspose2d(hid_channels*2, out_channels,\n                               kernel_size=4, stride=2, padding=1),\n            nn.Tanh(),\n        )\n\n        # Apply Init Weights\n        # Note Downsampling and Upsampling already applied\n        self.feature_block.apply(init_weights)\n\n       \n    # Establish Skip Connections in a symmetrical fashion\n    def forward(self, x):\n\n        # Begin by applying all downsamples while storing the order\n        skips = []\n        for down in self.downsampling_path:\n            x = down(x)\n            skips.append(x)\n        skips = reversed(skips[:-1])\n\n        # Next apply all upsamples then concat result with\n        # Corresponding downsample result before again upsampling and repeat\n        for up, skip in zip(self.upsampling_path, skips):\n            x = up(x)\n            x = torch.cat([x, skip], dim=1)\n\n        return self.feature_block(x)","metadata":{"id":"mSsuefzpdDXG","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Quick Sanity Check to Ensure the UNet architecture operates on our image size correctly, generating a new image of the same size","metadata":{"id":"lkTItOEAKMat"}},{"cell_type":"code","source":"def test_UNet():\n  # Test with input image of size 3 x 256 x 256 with batch_size = 5\n  in_channels = 3\n  hidden_channels = 64\n  out_channels = 3\n  batch_size = 5\n  H = 256\n  W = 256\n  # Create Zero Tensor of batch x image size\n  image = torch.zeros(((batch_size, in_channels, H, W)))\n  print(\"Initial Size:\")\n  print(image.shape)\n  print()\n\n  # Apply UNet Generator!\n  model = UNetGenerator(hidden_channels, in_channels, out_channels)\n  generated_image = model(image)\n  # You should see that UNet ends with the same size image generated\n  # In this Case [batch_size, out_channels, H, W]\n  print(\"UNet Generation Size:\")\n  print(generated_image.size())\n\ntest_UNet()","metadata":{"executionInfo":{"elapsed":4069,"status":"ok","timestamp":1684109608592,"user":{"displayName":"Garrett Devereux","userId":"06903598260790252274"},"user_tz":420},"id":"7tZqDUCMKME8","outputId":"0c27d08c-00d7-4a9c-a0be-d91de8ca0e0a","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Discriminator\n\nThe discriminator takes in the input image and classifies it as real or fake (generated). However, instead of conventially outputing a single probability of the input image being real or fake, our GAN used the [PatchGAN discriminator](https://paperswithcode.com/method/patchgan). \"PatchGAN is a type of discriminator which only penalizes structure at the scale of local image patches. The PatchGAN discriminator tries to classify if each patch in an image is real or fake. This discriminator is run convolutionally across the image, averaging all responses to provide the ultimate output decision.\" \n\nThis implementation will allow more descriptive losses that will help us improve our generators faster than a single Real or Fake prediction might.\n\n![image2.jpg](https://www.researchgate.net/profile/Gozde-Unal-2/publication/323904616/figure/fig1/AS:606457334595585@1521602104652/PatchGAN-discriminator-Each-value-of-the-output-matrix-represents-the-probability-of_W640.jpg)\n\n[Image Source](https://www.researchgate.net/publication/323904616_Patch-Based_Image_Inpainting_with_Generative_Adversarial_Networks)","metadata":{"id":"Mbbo3Z-_fHq8"}},{"cell_type":"code","source":"class Discriminator(nn.Module):\n    def __init__(self, hid_channels, in_channels=3):\n        super().__init__()\n\n        # Initialize Series of Downsamples and Convolutions to map\n        # (batch_size, 3, 256, 256) -> (batch_size, 512, 1, 1)\n        self.block = nn.Sequential(\n            # Note all following have batch_size as 0 dimension\n            # And are assuming hid=64, in=3, H=W=256\n            # However this function will generalize for other sizes\n            # (3, 256, 256) -> (64, 128, 128)\n            Downsampling(in_channels, hid_channels, norm=False),\n            # (64, 128, 128) -> (128, 64, 64)\n            Downsampling(hid_channels, hid_channels*2),\n            # (128, 64, 64) -> (256, 32, 32)\n            Downsampling(hid_channels*2, hid_channels*4),\n            # Note Stride 1 so H, W do not half\n            # (256, 32, 32) -> (512, 31, 31)\n            Downsampling(hid_channels*4, hid_channels*8, stride=1),\n            # Final Convolution mapping\n            # (512, 31, 31) -> (1, 30, 30)\n            nn.Conv2d(hid_channels*8, 1, kernel_size=4, padding=1), # 1x30x30\n        )\n\n        # Apply Init Weights\n        self.block.apply(init_weights)\n        \n    def forward(self, x):\n        return self.block(x)","metadata":{"id":"b7ecrbRJdzN1","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Quick Sanity Check to Ensure the PatchGAN discriminator operates on our image size correctly, generating a matrix of probability scores of size 1 x 30 x 30. Note that the probability scores here are meaningless as the discriminator has not yet been trained","metadata":{"id":"Z2Sokz_MUJYe"}},{"cell_type":"code","source":"def test_Patch():\n  # Test with input image of size 3 x 256 x 256 with batch_size = 5\n  in_channels = 3\n  hidden_channels = 64\n  out_channels = 3\n  batch_size = 5\n  H = 256\n  W = 256\n  # Create Zero Tensor of batch x image size\n  image = torch.ones(((batch_size, in_channels, H, W)))\n  print(\"Initial Size:\")\n  print(image.shape)\n  print()\n\n  # Apply Patch Discriminator!\n  model = Discriminator(hidden_channels, in_channels)\n  probability_scores = model(image)\n  # You should see that Discriminator ends with \n  # probability matrix of size  [batch_size, 1, 30, 30]\n  print(\"PatchGan Probability Scores Size:\")\n  print(probability_scores.size())\n  print()\n\n  print(\"PatchGan Probability Scores:\")\n  print(probability_scores)\n\ntest_Patch()","metadata":{"executionInfo":{"elapsed":738,"status":"ok","timestamp":1684109626167,"user":{"displayName":"Garrett Devereux","userId":"06903598260790252274"},"user_tz":420},"id":"b2XCBmpoUGdl","outputId":"340a141b-b9f5-48bd-9395-f7f466115cdc","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Visualizing the Generators and Discriminators - Before Training\n\nNow we can initialize untrained Generators and Discriminators. Note that nothing has been trained yet, so the behavior is meaningless at this point","metadata":{"id":"fEW50f9ceNSP"}},{"cell_type":"code","source":"# Base Dimensions to use\nhid_channels = 64\nin_channels = 3\nout_channels = 3\n\n# Generator to transform Photos to Monet Paintings\nmonet_generator = UNetGenerator(hid_channels, in_channels, out_channels)\n# Generator to transform Monet Paintings to Photos\nphoto_generator = UNetGenerator(hid_channels, in_channels, out_channels)\n\n# Give a batch of example photos to each generator\nto_monet = monet_generator(example_photos)\nto_photos = photo_generator(example_monets)\n\nshow_img(example_photos, nrow=5, title=\"Original Photos\")\nshow_img(to_monet, nrow=5, title=\"Generated Monets\")\nshow_img(example_monets, nrow=5, title=\"Original Monets\")\nshow_img(to_photos, nrow=5, title=\"Generated Photos\")","metadata":{"executionInfo":{"elapsed":22353,"status":"ok","timestamp":1684109653977,"user":{"displayName":"Garrett Devereux","userId":"06903598260790252274"},"user_tz":420},"id":"3LvRnlEVeaPW","outputId":"eded2845-2aac-4294-a4e1-4c132ce8a5fd","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"We can do the same with Discriminators, passing in photos to see the probability scores of the image being generated or from the Dataset. These have also not been trained yet and have meaninless results.","metadata":{"id":"Ndybvs_kfiWZ"}},{"cell_type":"code","source":"# Discriminator to identify real Monet Paintings\nmonet_discriminator = Discriminator(hid_channels)\n# Discriminator to identify real photos\nphoto_discriminator = Discriminator(hid_channels)\n\n# Give a real photo to each discriminator (first one photo shown above)\nreal_monet_scores = monet_discriminator(example_monets[0])\nreal_photo_scores = photo_discriminator(example_photos[0])\n\n# Sizes are randomly initialized\nprint(\"Monet Discriminator Size and Scores\")\nprint(real_monet_scores.shape)\nprint(real_monet_scores)\nprint()\nprint(\"Photo Discriminator Size and Scores\")\nprint(real_photo_scores.shape)\nprint(real_photo_scores)","metadata":{"executionInfo":{"elapsed":437,"status":"ok","timestamp":1684109665575,"user":{"displayName":"Garrett Devereux","userId":"06903598260790252274"},"user_tz":420},"id":"EqXpm9eaf0UY","outputId":"161f1c23-a8d7-466d-973b-69f9bbfa6f06","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Build THE CycleGAN!\n\nWe now have our Generator and Discriminator and are ready to build the CylcleGAN. First Let's explore the architecture and loss functions:\n\n","metadata":{"id":"8HhJGCR4gXrl"}},{"cell_type":"markdown","source":"### Loss Functions\n\nThe training process consists of competing Generators and Discriminator with a variety of loss functions. To simplify the computation graph, we will explore each section one at a time and work to specify adversarial loss, identity loss, cycle loss, generator loss, and discriminator loss. Note that we are using MSE loss for any comparison with discriminator scores, and L1 loss for comparisons between images. These loss functions may be changed in to future to increase model performance.","metadata":{"id":"yAw5sa_IoAVy"}},{"cell_type":"markdown","source":"### Generator Loss\n\nOur Generator Loss is comprised of three loss functions, each with their own objective:","metadata":{"id":"PPjVeQXAhlim"}},{"cell_type":"markdown","source":"1. **Adversarial Loss**: This loss evaluates how well a generator is tricking a discriminator. With a perfect generator, our fake generated image would produce a perfect discriminator score -> a matrix of all 1s. This loss function simply takes in the discriminator score of a fake image and compares it to a corresponding tensor of all 1s, evaluating how well we are currently tricking the discriminator\n\n![image3.png](https://raw.githubusercontent.com/garrettdevereux/DeepLearningGenerativeArt/main/AdvLoss.png)","metadata":{"id":"BgWHijynkVtB"}},{"cell_type":"code","source":"# How good are we at tricking the discriminator\ndef get_adv_loss(fake_image_discriminator_score):\n    # A perfect fake will have the discriminator produce all 1s\n    real_labels = torch.ones_like(fake_image_discriminator_score)\n    # Retrun mse between actual scores and 1s\n    return F.mse_loss(fake_image_discriminator_score, real_labels)","metadata":{"id":"YLh8EtJOktSQ","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"2. **Cycle Loss**: This loss evaluates how similar our twice transformed images are. Part of the CycleGAN training process has us take a real photo from the dataset, use the first generator to produce a fake generated version of the image, and then feed this new fake photo back through the second generator to return back to the original domain. With perfect generators, this should be the same image as the original. This loss simply compares the L1 distance between the original image and the reconstructed image.\n\n![image4.png](https://raw.githubusercontent.com/garrettdevereux/DeepLearningGenerativeArt/main/CycleLoss.png)","metadata":{"id":"xLcw6v_mk60k"}},{"cell_type":"code","source":"# How similar are our twice transformed photos\n# i.e real -> generated -> recon should have similar real and recon\ndef get_cycle_loss(real, recon, lambda_cycle):\n    # Simply return reg * L1 loss between recon image and real image\n    return lambda_cycle * F.l1_loss(recon, real)","metadata":{"id":"vw9wpjyZlkpx","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"3. **Identity Loss**: This loss evaluates how similar the original real image is compared with its identity image. Another part of the CycleGAN training step is passing a real image (such as a real monet paiting) into the monet generator whose goal is to produce monet paintings. As the input is already in the correct domain, a perfect generator would not change the image at all. This loss simply compares the L1 distance between the original real image and the generated identity image.\n\n![image5.png](https://raw.githubusercontent.com/garrettdevereux/DeepLearningGenerativeArt/main/IdentityLoss.png)","metadata":{"id":"DvDK41e8lqVF"}},{"cell_type":"code","source":"# How similar the real and identity images are\n# Given a photo in the correct domain, our generator should do very little\ndef get_idt_loss(real, identity_image, lambda_cycle, lambda_idt):\n    # Simply return reg * L1 loss between real and identity\n    return lambda_idt * lambda_cycle * F.l1_loss(identity_image, real)","metadata":{"id":"PBxx6ehUmlYH","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Total Generator Loss\nFinally, we put everything together to get the total loss of BOTH generators. This involves summing the total adversarial loss, cycle loss, and identity loss of both generators and discriminators:\n\n![image6.png](https://raw.githubusercontent.com/garrettdevereux/DeepLearningGenerativeArt/main/GeneratorLoss.png)","metadata":{"id":"ccC1HF-imojP"}},{"cell_type":"code","source":"# Computes the Loss of BOTH Generators, given\n# - Original Photos from dataset, real_monet and real_photo\n# - Generated Fakes, fake_monet and fake_photo\n# - Reconstructed Photos, reconstructed_monet, reconstructed_photo\n# - Identities, monet_identity and photo_identity\n# - Discrim Scores of fake images, fake_monet_discrim_score, fake_photo_discrim_score\n# - Model Hyperparams lambda_cycle and lambda_idt\ndef get_gen_loss(real_monet, real_photo, fake_monet, fake_photo, reconstructed_monet, reconstructed_photo,\n                 monet_identity, photo_identity, fake_monet_discrim_score, fake_photo_discrim_score, lambda_cycle, lambda_idt):\n  \n\n    # STEP 1: Compute Adversarial Loss using fakes and discriminator scores\n    # Conceptually - how good the generatros are at tricking the discriminator\n    # Adv Loss of monet generator\n    adv_loss_photo_to_monet = get_adv_loss(fake_monet_discrim_score)\n    # Adv Loss of photo generator\n    adv_loss_monet_to_photo = get_adv_loss(fake_photo_discrim_score)\n    total_adv_loss = adv_loss_photo_to_monet + adv_loss_monet_to_photo\n            \n    # STEP 2: Compute Cycle Loss using\n    # Cycle loss of the original monet compared to reconstructed monet\n    cycle_loss_MPM = get_cycle_loss(real_monet, reconstructed_monet, lambda_cycle[0])\n    # Cycle loss of the original photo compared to reconstructed photo\n    cycle_loss_PMP = get_cycle_loss(real_photo, reconstructed_photo, lambda_cycle[1])\n    total_cycle_loss = cycle_loss_MPM + cycle_loss_PMP\n\n    # STEP 3: Compute Identity Loss using reals and identities\n    # Identity Loss of the real monet and monet identity\n    idt_loss_MM = get_idt_loss(real_monet, monet_identity, lambda_cycle[0], lambda_idt)\n    # Identity Loss of the real photo and identity photo\n    idt_loss_PP = get_idt_loss(real_photo, photo_identity, lambda_cycle[1], lambda_idt)\n    total_idt_loss = idt_loss_MM + idt_loss_PP\n            \n    # FINAL STEP: Sum losses to get total generator loss\n    gen_loss = total_adv_loss + total_idt_loss + total_cycle_loss\n    return gen_loss","metadata":{"id":"lkIdtRoDueDy","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Discriminator Loss\n\nIn contrast to the Generator Loss that computes the loss of BOTH Generators, our Discriminator loss is simpler and only computed the loss of a single Discriminator. This is very similar to the Adversarial Loss, but now is evaluating how well the discriminator is performing on real and generated images. With a perfect Discriminator, our fake generated image would produce a discriminator score of all 0s, while a real image produces a discriminator score of all 1s. This loss function simply takes in the discriminator score of the real image and compares it with the MSE loss of all 1s, and then uses the discriminator score of a fake image and computed the MSE loss with a tensor of all 0s. Finally, it averages these scores to produce the final discriminator loss\n\n![image7.png](https://raw.githubusercontent.com/garrettdevereux/DeepLearningGenerativeArt/main/DiscriminatorLoss.png)","metadata":{"id":"z0HIvxSzoCz9"}},{"cell_type":"code","source":"# Computes loss of Single Discriminator given fake and real discriminator scores\ndef get_disc_loss(real_discrim_score, fake_discrim_score):\n    # STEP 1: Calculate real discrim loss\n    # Our discriminator should output all 1s as this is a real image from the dataset\n    real_labels = torch.ones_like(real_discrim_score)\n    # Return mse between actual scores and 1s\n    real_loss = F.mse_loss(real_discrim_score, real_labels)\n        \n    # STEP 2: Calculate fake discrim loss\n    # Our discriminator should output all 0s as this is a generated image from the dataset\n    fake_labels = torch.zeros_like(fake_discrim_score)\n    # Return mse between actual scores and 1s\n    fake_loss = F.mse_loss(fake_discrim_score, fake_labels)\n        \n    # Return average of losses\n    disc_loss = (fake_loss + real_loss) * 0.5\n    return disc_loss","metadata":{"id":"CZMXvaBmn9Sg","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### PyTorch Lightning Module\nAs we are using PyTorch Lightning, our CycleGAN will be a ```LightningModule```. This takes helps simplify the training process taking care of what device the code is run on while organizing our code into six parts:\n\n*   Initialization (```__init__```).\n*   Train Loop (```training_step()```)\n*   Validation Loop (```validation_step()```)\n*   Test Loop (```test_step()```)\n*   Prediction Loop (```predict_step()```)\n*   Optimizers and LR Schedulers (```configure_optimizers()```)\n\nSee the [PyTorch LightningModule Documentation](https://lightning.ai/docs/pytorch/stable/common/lightning_module.html) for more information. Also, note for the base model we are not using and LR Schedulers or decay, this can be added for more complexity later","metadata":{"id":"TzLyf469m6fh"}},{"cell_type":"code","source":"class CycleGAN(L.LightningModule):\n    def __init__(\n        self,\n        in_channels,\n        out_channels,\n        hid_channels,\n        optimizer,\n        lr,\n        betas,\n        lambda_idt,\n        lambda_cycle,\n        num_epochs,\n        monet_gen = None,\n        photo_gen = None,\n        monet_disc = None,\n        photo_disc = None\n    ):\n        super().__init__()\n        self.save_hyperparameters(ignore=[\"optimizer\"])\n        self.optimizer = optimizer\n        # Necessary to set as we are using multiple optimizers\n        self.automatic_optimization = False\n\n        # Step for Graphs\n        self.step_num = 0\n        \n        # INITIALIZE OUR GENERATORS AND DISCRIMINATORS!!!!\n        # Weights will automatically be initialized here\n\n        # Generators:\n        # Generator to transform Photos to Monet Paintings\n        if monet_gen == None:\n          self.monet_generator = UNetGenerator(hid_channels, in_channels, out_channels)\n        else:\n          self.monet_generator = monet_gen\n  \n        # Generator to transform Monet Paintings into Photos\n        if photo_gen == None:\n          self.photo_generator = UNetGenerator(hid_channels, in_channels, out_channels)\n        else:\n          self.photo_generator = photo_gen\n\n        # Discriminators\n        # Discriminator to identify real Monet Paintings\n        if monet_disc == None:\n          self.monet_discriminator = Discriminator(hid_channels)\n        else:\n          self.monet_discriminator = monet_disc\n\n        # Discriminator to identify real Photos\n        if photo_disc == None:\n          self.photo_discriminator = Discriminator(hid_channels)\n        else:\n          self.photo_discriminator = photo_disc\n\n        # Metrics for Loss graphs\n        self.losses = {\n            \"gen_loss\": [],\n            \"disc_loss_M\": [],\n            \"disc_loss_P\": [],\n        }\n        \n    # As our predominant focus is creating the best Monet Painting Generator\n    # Our forward simply passes in an image to the Monet Generator\n    def forward(self, img):\n        return self.monet_generator(img)   \n    \n\n    # Configure the Optimizers for both our generator and discriminator\n    def configure_optimizers(self):\n        # Load hyperparams\n        opt_config = {\n            \"lr\": self.hparams.lr,\n            \"betas\": self.hparams.betas,\n        }\n        # Set up optimizer for generators\n        opt_gen = self.optimizer(\n            list(self.monet_generator.parameters()) + list(self.photo_generator.parameters()),\n            **opt_config,\n        )\n        # Set up optimizer for discriminators\n        opt_disc = self.optimizer(\n            list(self.monet_discriminator.parameters()) + list(self.photo_discriminator.parameters()),\n            **opt_config,\n        )\n        optimizers = [opt_gen, opt_disc]\n        # No schedulers for now\n        schedulers = []\n        return optimizers, schedulers\n    \n    def training_step(self, batch, batch_idx):\n        # Begin by storing batches of real monets and photos from the dataset\n        real_monet = batch[\"monet\"]\n        real_photo = batch[\"photo\"]\n\n        # Store the optimizers for the generators and discriminators\n        gen_optimizer, disc_optimizer = self.optimizers()\n\n        # STEP 1: Use Generators to produce fake images in both domains\n        # Fake monet generated from the real photo\n        fake_monet = self.monet_generator(real_photo)\n        # Fake photo generated from the real monet\n        fake_photo = self.photo_generator(real_monet)\n         \n        # STEP 2: Use generators and fakes to reconstruct original inputs\n        # Reconstruction of original monet painting\n        reconstructed_monet = self.monet_generator(fake_photo)\n        # Reconstruction of original photo\n        reconstructed_photo = self.photo_generator(fake_monet)\n\n        # STEP 3: Construct identities by feeding real inputs into the generators\n        # Monet identity: feed real monet into monet generator\n        monet_identity = self.monet_generator(real_monet)\n        # Photo identity: feed real photo into photo generator\n        photo_identity = self.photo_generator(real_photo)\n\n        # STEP 4: Compute Discriminator Scores to see how well our\n        # Discriminators differentiate real and fake images\n        # Scores of how good the fake generated monets are\n        fake_monet_discrim_score = self.monet_discriminator(fake_monet)\n        # Scores of how good the fake generated photos are\n        fake_photo_discrim_score = self.photo_discriminator(fake_photo)\n        # Scores of how good the real monets are\n        real_monet_discrim_score = self.monet_discriminator(real_monet)\n        # Scores of how good the real photos are\n        real_photo_discrim_score = self.photo_discriminator(real_photo)\n    \n        # STEP 5: Use Generator Loss to train Generators\n        # Standard pipeline: Loss -> zero_grad -> backward -> step\n        self.toggle_optimizer(gen_optimizer, 0)\n        gen_loss = get_gen_loss(real_monet, real_photo, fake_monet, fake_photo, reconstructed_monet, reconstructed_photo,\n                 monet_identity, photo_identity, fake_monet_discrim_score, fake_photo_discrim_score, self.hparams.lambda_cycle, self.hparams.lambda_idt)      \n        gen_optimizer.zero_grad()\n        self.manual_backward(gen_loss)\n        gen_optimizer.step()\n        self.untoggle_optimizer(0)\n        \n        # STEP 5: Use Discriminator Loss to train Discriminators\n        # Standard pipeline: Losses -> zero_grad -> backward -> step\n        self.toggle_optimizer(disc_optimizer, 1)\n        # Need to detatch otherwise backward is happening twice\n        monet_disc_loss = get_disc_loss(real_monet_discrim_score, self.monet_discriminator(fake_monet.detach()))\n        photo_disc_loss = get_disc_loss(real_photo_discrim_score, self.photo_discriminator(fake_photo.detach()))\n        disc_optimizer.zero_grad()\n        self.manual_backward(monet_disc_loss)\n        self.manual_backward(photo_disc_loss)\n        disc_optimizer.step()\n        self.untoggle_optimizer(1)\n        \n        # Done with training step\n        # Log metrics and losses\n        metrics = {\n            \"gen_loss\": gen_loss,\n            \"disc_loss_M\": monet_disc_loss,\n            \"disc_loss_P\": photo_disc_loss,\n        }\n        self.log_dict(metrics, on_step=False, on_epoch=True, prog_bar=True)\n\n        # Add to Losses 4 times each epoch\n        if (self.step % 1757 == 0):\n          self.losses[\"gen_loss\"].append(float(gen_loss))\n          self.losses[\"disc_loss_M\"].append(float(monet_disc_loss))\n          self.losses[\"disc_loss_P\"].append(float(photo_disc_loss))\n\n        self.step += 1\n        \n        \n    # Output current progress of generator on a validation set\n    def validation_step(self, batch, batch_idx):\n        self.display_results(batch, batch_idx, \"validate\")\n    \n    # Output current progress of generator on a test set\n    def test_step(self, batch, batch_idx):\n        self.display_results(batch, batch_idx, \"test\")\n        \n    # Predict the results of a batch\n    def predict_step(self, batch, batch_idx):\n        return self(batch)\n    \n    # Function to Visualize Results\n    def display_results(self, batch, batch_idx, stage):\n        # Store batch of real input photos and corresponding fake generated monets\n        real_photos = batch\n        fake_monets = self(real_photos)\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        # Display current image results\n        show_img(\n            torch.cat([real_photos, fake_monets], dim=0),\n            nrow=len(real_photos),\n            title=title,\n        )\n    \n    def on_train_epoch_start(self):\n        # Update Progress Bar\n        self.log(\"lr\", self.hparams.lr, on_step=False, on_epoch=True, prog_bar=True)\n\n        # Reset Step\n        self.step = 0\n        \n    def on_train_epoch_end(self):\n        # Print Current State of Models\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        \"\"\"\n        print(\"Saving Final Models.\")\n        # Save Models into Checkpoints\n        monet_gen_path = checkpoints + \"/monet_generator\" + str(self.current_epoch) + \".pth\"\n        torch.save(self.monet_generator.state_dict(), monet_gen_path)\n\n        photo_gen_path = checkpoints + \"/photo_generator\" + str(self.current_epoch) + \".pth\"\n        torch.save(self.photo_generator.state_dict(), photo_gen_path)\n\n        monet_disc_path = checkpoints + \"/monet_discriminator\" + str(self.current_epoch ) + \".pth\"\n        torch.save(self.monet_discriminator.state_dict(), monet_disc_path)\n\n        photo_disc_path = checkpoints + \"/photo_discriminator\" + str(self.current_epoch) + \".pth\"\n        torch.save(self.photo_discriminator.state_dict(), photo_disc_path)\n        print(\"Training ended.\")\n        \"\"\"\n        \n    def on_predict_epoch_end(self, results):\n        predictions = self.trainer.predict_loop.predictions\n        print(f\"Done predicting\")","metadata":{"id":"VIZx1_P6hreW","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Training Configuration\n\nWe use PyTorch Lightning Training Object to do everything for us! Here we set up the requuired Configuration and Params.","metadata":{"id":"0e3hp-ynjc-K"}},{"cell_type":"code","source":"MODEL_PARAMS = {    \n    # In channels and out channels\n    \"in_channels\": 3,\n    \"out_channels\": 3,\n\n    # the number of filters in the first layer for the generators and discriminators\n    \"hid_channels\": 64,\n    \n    # using Adam\n    \"optimizer\": 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 number of epochs for training\n    \"num_epochs\": 8,\n}","metadata":{"id":"QT_DFMILg-aS","trusted":true},"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,\n    \n    # train on single GPU\n    \"devices\": 1,\n    \n    # We are manually checkpointing inside fo the GAN\n    \"enable_checkpointing\": False,\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_PARAMS[\"num_epochs\"],\n    \"limit_train_batches\": 1.0,\n    \"limit_predict_batches\": 1.0,\n    \n    # the maximum amount of time for training, in case we exceed run-time of 5 hours\n    \"max_time\": {\"hours\": 11, \"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,\n}","metadata":{"id":"JnVIXbjtjmTQ","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## TRAIN THE CYLCEGAN\nWe have now defined all required elements in order to train our model. We now use PyTorch Lightning to do the training, instantiating a ```CustomDataModule``` object to handle the Dataset, our ```CylceGAN``` to store the model parameters and controlling all of the steps, and a PyTorch ```Trainer``` object that handles the entire training process. Calling ```trainer.fit``` with the model and datamodule will then train the model.","metadata":{"id":"xIALLEsDMvz3"}},{"cell_type":"code","source":"# Initialize Trainer with model\ndm = CustomDataModule(**DM_CONFIG)\nmodel = CycleGAN(**MODEL_PARAMS)\ntrainer = L.Trainer(**TRAIN_CONFIG)\n\ntrainer.fit(model, datamodule=dm)","metadata":{"id":"ue4KGJpvj2_x","outputId":"4ca1363f-fd15-4751-b966-679ae90cc66d","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Post Training Visualization\n\nLet's visualize how the Loss of both of the Generators as well as the individual Discriminator Losses changed throughout the training process.","metadata":{"id":"Zi6IjSqZqSb_"}},{"cell_type":"code","source":"# Loss Plots\nfig, axes = plt.subplots(ncols=1, nrows=2, figsize=(18, 12), facecolor='w')\ntraining_steps = np.arange(len(model.losses[\"gen_loss\"])) / 4\n\n# Plot the Generator Loss\naxes[0].set_title(\"Generator Loss\")\naxes[0].plot(training_steps, model.losses[\"gen_loss\"], label='Generator Loss')\naxes[0].legend()\naxes[0].set_xlabel('Epochs')\n\n# Plot the Discriminator Losses\naxes[1].set_title(\"Discriminator Losses\")\naxes[1].plot(training_steps, model.losses[\"disc_loss_M\"], label='Monet Discriminator Loss')\naxes[1].plot(training_steps, model.losses[\"disc_loss_P\"], label='Photo Discriminator Loss')\naxes[1].legend()\naxes[1].set_xlabel('Epochs')\n\nplt.show()","metadata":{"executionInfo":{"elapsed":1713,"status":"ok","timestamp":1683872676459,"user":{"displayName":"Garrett Devereux","userId":"06903598260790252274"},"user_tz":420},"id":"VJE4FwVGF36B","outputId":"adb2d61d-0b69-4157-eea3-9eba2dc58810","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Visualizing Trained Generators and Discriminators\nLet's now feed batches of images of photos and monet paintings to our generators to see how they perform. On top of generated fakes, let's also generate the cyclic and identity photos. Here we also print out Discriminator Results. Analysis of results","metadata":{"id":"ANLmKqlYsbK6"}},{"cell_type":"code","source":"# Contruct Fakes, Reconstructions, and Identities \n# with a batch of example photos and money paintings\n\nexample_photos = example_photos\nexample_monets = example_monets\n\nfake_monets = model.monet_generator(example_photos)\nfake_photos = model.photo_generator(example_monets)\n\nrecon_monet = model.monet_generator(fake_photos)\nrecon_photo = model.photo_generator(fake_monets)\n\nidt_monet = model.monet_generator(example_monets)\nidt_photo = model.photo_generator(example_photos)\n\n# Discriminator Scores\nreal_monet_score = model.monet_discriminator(example_monets[0])\nreal_photo_score = model.photo_discriminator(example_photos[0])\nfake_monet_score = model.monet_discriminator(fake_monets[0])\nfake_photo_score = model.photo_discriminator(fake_photos[0])\nones_label = torch.ones_like(real_monet_score)\nzeros_label = torch.zeros_like(real_monet_score)\n\nprint(\"Visualizing Results produced from Original Photos\")\nprint(\"Real Photo Discriminator Score (on first image)\")\nprint(real_photo_score)\nprint(\"Real Photo Loss\")\nprint(F.mse_loss(real_photo_score, ones_label))\n# Visualize from Original Photos\nshow_img(example_photos, nrow=5, title=\"Original Photos\")\n\nshow_img(fake_monets, nrow=5, title=\"Generated Monets\")\nprint(\"Fake Monet Discriminator Score (on first image)\")\nprint(fake_monet_score)\nprint(\"Fake Monet Loss\")\nprint(F.mse_loss(fake_monet_score, zeros_label))\nprint()\n\nshow_img(recon_photo, nrow=5, title=\"Reconstructed Photos\")\n\nshow_img(idt_photo, nrow=5, title=\"Identity Photos\")\n\n# Visualize from Original Monets\nprint(\"Visualizing Results produced from Original Monets\")\nprint(\"Real Monet Discriminator Score (on first image)\")\nprint(real_monet_score)\nprint(\"Real Monet Loss\")\nprint(F.mse_loss(real_monet_score, ones_label))\nprint()\nshow_img(example_monets, nrow=5, title=\"Original Monets\")\n\nshow_img(fake_photos, nrow=5, title=\"Generated Photos\")\nprint(\"Fake Photo Discriminator Score (on first image)\")\nprint(fake_photo_score)\nprint(\"Fake Photo Loss\")\nprint(F.mse_loss(fake_photo_score, zeros_label))\nprint()\n\nshow_img(recon_monet, nrow=5, title=\"Reconstructed Monets\")\n\nshow_img(idt_monet, nrow=5, title=\"Identity Monets\")","metadata":{"executionInfo":{"elapsed":57356,"status":"ok","timestamp":1684110025023,"user":{"displayName":"Garrett Devereux","userId":"06903598260790252274"},"user_tz":420},"id":"y0OqwP7xrNfA","outputId":"10444ba0-dfac-4c59-8286-47c427226c86","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Submit To Kaggle\n\nFinally, we can submit our final model to Kaggle!","metadata":{"id":"yQohDU2Jy7Ds"}},{"cell_type":"code","source":"# Generate all monet paintings from training photos\npredictions = trainer.predict(model, datamodule=dm)","metadata":{"id":"OHy8Cno-zgAx","executionInfo":{"status":"ok","timestamp":1684111013494,"user_tz":420,"elapsed":71846,"user":{"displayName":"Garrett Devereux","userId":"06903598260790252274"}},"outputId":"94d579e5-6100-46bc-cf7c-c9fd43379a2b","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"os.makedirs(\"../images\", exist_ok=True)\n\nidx = 0\nfor img in predictions:\n    # Extract 3x256x256 image\n    pred_img = img[0]\n    # Save image to desired directory\n    save_image(\n            pred_img.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\")\nprint(\"Submission Complete\")","metadata":{"id":"YKjieWge1to0","executionInfo":{"status":"ok","timestamp":1684111807362,"user_tz":420,"elapsed":64572,"user":{"displayName":"Garrett Devereux","userId":"06903598260790252274"}},"outputId":"8fafea10-bcdd-4706-9040-262e87564a06","execution":{"iopub.status.busy":"2023-05-26T13:49:23.063766Z","iopub.execute_input":"2023-05-26T13:49:23.064602Z","iopub.status.idle":"2023-05-26T13:49:23.479683Z","shell.execute_reply.started":"2023-05-26T13:49:23.064557Z","shell.execute_reply":"2023-05-26T13:49:23.476771Z"},"trusted":true},"execution_count":null,"outputs":[]}]}