{"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 one of our advanced implementations 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 implement Differentiable Augmentation for a Dual Discriminator and compare it to the performance of our [Baseline DC CycleGAN](https://github.com/garrettdevereux/DeepLearningGenerativeArt/blob/main/MonetGANV1.ipynb).\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* [Two-objective Discriminator Notebook](https://www.kaggle.com/code/unfriendlyai/two-objective-discriminator)\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* [Original Differential Augmentation paper](https://arxiv.org/abs/2006.10738)\n* [Original Dual Discriminator paper](https://arxiv.org/abs/1709.03831)\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;\n!pip install deepspeed","metadata":{"executionInfo":{"elapsed":45531,"status":"ok","timestamp":1685506226657,"user":{"displayName":"Garrett Devereux","userId":"06903598260790252274"},"user_tz":420},"id":"-mQnQb6rm1C7","outputId":"743a704b-46e7-4942-c26a-e0fd430f0e18","execution":{"iopub.status.busy":"2023-06-01T00:54:49.707467Z","iopub.execute_input":"2023-06-01T00:54:49.707861Z","iopub.status.idle":"2023-06-01T00:55:23.550462Z","shell.execute_reply.started":"2023-06-01T00:54:49.707832Z","shell.execute_reply":"2023-06-01T00:55:23.549541Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import glob\nimport os\nimport shutil\n\nimport deepspeed as ds\nimport torchvision.models as models\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\n\ndevice = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\nprint(device)","metadata":{"executionInfo":{"elapsed":4509,"status":"ok","timestamp":1685506231154,"user":{"displayName":"Garrett Devereux","userId":"06903598260790252274"},"user_tz":420},"id":"zXMDumZap7i4","outputId":"bbf8121d-46d8-4a44-86cd-c3a62d4ef3b8","execution":{"iopub.status.busy":"2023-06-01T00:55:23.553146Z","iopub.execute_input":"2023-06-01T00:55:23.553819Z","iopub.status.idle":"2023-06-01T00:55:36.420823Z","shell.execute_reply.started":"2023-06-01T00:55:23.553778Z","shell.execute_reply":"2023-06-01T00:55:36.419868Z"},"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","execution":{"iopub.status.busy":"2023-06-01T00:55:36.424780Z","iopub.execute_input":"2023-06-01T00:55:36.425692Z","iopub.status.idle":"2023-06-01T00:55:36.435592Z","shell.execute_reply.started":"2023-06-01T00:55:36.425657Z","shell.execute_reply":"2023-06-01T00:55:36.434528Z"},"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","execution":{"iopub.status.busy":"2023-06-01T00:55:36.441438Z","iopub.execute_input":"2023-06-01T00:55:36.442110Z","iopub.status.idle":"2023-06-01T00:55:36.459271Z","shell.execute_reply.started":"2023-06-01T00:55:36.442076Z","shell.execute_reply":"2023-06-01T00:55:36.458236Z"},"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","execution":{"iopub.status.busy":"2023-06-01T00:55:36.463817Z","iopub.execute_input":"2023-06-01T00:55:36.466707Z","iopub.status.idle":"2023-06-01T00:55:36.474646Z","shell.execute_reply.started":"2023-06-01T00:55:36.466672Z","shell.execute_reply":"2023-06-01T00:55:36.473764Z"},"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":53671,"status":"ok","timestamp":1685506315198,"user":{"displayName":"Garrett Devereux","userId":"06903598260790252274"},"user_tz":420},"id":"GnmReHHjQsZF","outputId":"c28045e1-1069-4ee2-afe5-12f866fcb191","execution":{"iopub.status.busy":"2023-06-01T00:55:36.479017Z","iopub.execute_input":"2023-06-01T00:55:36.480077Z","iopub.status.idle":"2023-06-01T00:55:43.554884Z","shell.execute_reply.started":"2023-06-01T00:55:36.480044Z","shell.execute_reply":"2023-06-01T00:55:43.553656Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Differential Augmentation\n\nFrom \"Differentiable Augmentation for Data-Efficient GAN Training\":\n\nThe performance of generative adversarial networks (GANs) heavily deteriorates\ngiven a limited amount of training data. This is mainly because the discriminator is memorizing the exact training set. To combat it, we propose Differentiable Augmentation (DiffAugment), a simple method that improves the data efficiency of GANs by imposing various types of differentiable augmentations on both real and fake samples. Previous attempts to directly augment the training data manipulate the distribution of real images, yielding little benefit; DiffAugment enables us to adopt the differentiable augmentation for the generated samples, effectively stabilizes training, and leads to better convergence.\"\n\nWith only 300 samples of Monet Paintings, Differential Augmentation seems to be a perfect fit for our model. The below implementation is inspired from the initial paper, and converts the TensorFlow code to PyTorch","metadata":{"id":"64C5SwKS6nbp"}},{"cell_type":"code","source":"# Inspired from https://github.com/mit-han-lab/data-efficient-gans/blob/master/DiffAugment-stylegan2/DiffAugment_tf.py\n\ndef DiffAugment(x, policy='', channels_first=True):\n    if policy:\n        if not channels_first:\n            x = x.permute(0, 3, 1, 2)\n        for p in policy.split(','):\n            for f in AUGMENT_FNS[p]:\n                x = f(x)\n        if not channels_first:\n            x = x.permute(0, 2, 3, 1)\n        x = x.contiguous()\n    return x\n\n\ndef rand_brightness(x):\n    x = x + (torch.rand(x.size(0), 1, 1, 1, dtype=x.dtype, device=x.device) - 0.5)\n    return x\n\n\ndef rand_saturation(x):\n    x_mean = x.mean(dim=1, keepdim=True)\n    x = (x - x_mean) * (torch.rand(x.size(0), 1, 1, 1, dtype=x.dtype, device=x.device) * 2) + x_mean\n    return x\n\n\ndef rand_contrast(x):\n    x_mean = x.mean(dim=[1, 2, 3], keepdim=True)\n    x = (x - x_mean) * (torch.rand(x.size(0), 1, 1, 1, dtype=x.dtype, device=x.device) + 0.5) + x_mean\n    return x\n\n\ndef rand_translation(x, ratio=0.125):\n    shift_x, shift_y = int(x.size(2) * ratio + 0.5), int(x.size(3) * ratio + 0.5)\n    translation_x = torch.randint(-shift_x, shift_x + 1, size=[x.size(0), 1, 1], device=x.device)\n    translation_y = torch.randint(-shift_y, shift_y + 1, size=[x.size(0), 1, 1], device=x.device)\n    grid_batch, grid_x, grid_y = torch.meshgrid(\n        torch.arange(x.size(0), dtype=torch.long, device=x.device),\n        torch.arange(x.size(2), dtype=torch.long, device=x.device),\n        torch.arange(x.size(3), dtype=torch.long, device=x.device),\n    )\n    grid_x = torch.clamp(grid_x + translation_x + 1, 0, x.size(2) + 1)\n    grid_y = torch.clamp(grid_y + translation_y + 1, 0, x.size(3) + 1)\n    x_pad = F.pad(x, [1, 1, 1, 1, 0, 0, 0, 0])\n    x = x_pad.permute(0, 2, 3, 1).contiguous()[grid_batch, grid_x, grid_y].permute(0, 3, 1, 2).contiguous()\n    return x\n\n\ndef rand_cutout(x, ratio=0.5):\n    cutout_size = int(x.size(2) * ratio + 0.5), int(x.size(3) * ratio + 0.5)\n    offset_x = torch.randint(0, x.size(2) + (1 - cutout_size[0] % 2), size=[x.size(0), 1, 1], device=x.device)\n    offset_y = torch.randint(0, x.size(3) + (1 - cutout_size[1] % 2), size=[x.size(0), 1, 1], device=x.device)\n    grid_batch, grid_x, grid_y = torch.meshgrid(\n        torch.arange(x.size(0), dtype=torch.long, device=x.device),\n        torch.arange(cutout_size[0], dtype=torch.long, device=x.device),\n        torch.arange(cutout_size[1], dtype=torch.long, device=x.device),\n    )\n    grid_x = torch.clamp(grid_x + offset_x - cutout_size[0] // 2, min=0, max=x.size(2) - 1)\n    grid_y = torch.clamp(grid_y + offset_y - cutout_size[1] // 2, min=0, max=x.size(3) - 1)\n    mask = torch.ones(x.size(0), x.size(2), x.size(3), dtype=x.dtype, device=x.device)\n    mask[grid_batch, grid_x, grid_y] = 0\n    x = x * mask.unsqueeze(1)\n    return x\n\n\nAUGMENT_FNS = {\n    'color': [rand_brightness, rand_saturation, rand_contrast],\n    'translation': [rand_translation],\n    'cutout': [rand_cutout],\n}\n\n# Should be passing a batch to this function (n, c, h, w)\ndef aug_fn(image):\n    return DiffAugment(image,\"color,translation,cutout\")","metadata":{"id":"belZCIWC72jW","execution":{"iopub.status.busy":"2023-06-01T00:55:43.556897Z","iopub.execute_input":"2023-06-01T00:55:43.557559Z","iopub.status.idle":"2023-06-01T00:55:43.580802Z","shell.execute_reply.started":"2023-06-01T00:55:43.557515Z","shell.execute_reply":"2023-06-01T00:55:43.579767Z"},"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, first with no augmentation:","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":35871,"status":"ok","timestamp":1685506351022,"user":{"displayName":"Garrett Devereux","userId":"06903598260790252274"},"user_tz":420},"id":"SecQv3LeRnb-","outputId":"bcfef35d-ac09-42f8-f050-ecb6e12d9ee0","execution":{"iopub.status.busy":"2023-06-01T00:55:43.582085Z","iopub.execute_input":"2023-06-01T00:55:43.582621Z","iopub.status.idle":"2023-06-01T00:55:44.586719Z","shell.execute_reply.started":"2023-06-01T00:55:43.582587Z","shell.execute_reply":"2023-06-01T00:55:44.585878Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Visualize with Differential Augmentation","metadata":{"id":"R5j6kwzgJiHB"}},{"cell_type":"code","source":"show_img(rand_brightness(example_monets), nrow=5, title=\"Random Brightness\")\nshow_img(rand_saturation(example_monets), nrow=5, title=\"Random Saturation\")\nshow_img(rand_contrast(example_monets), nrow=5, title=\"Random Contrast\")\nshow_img(rand_translation(example_monets), nrow=5, title=\"Random Translation\")\nshow_img(rand_cutout(example_monets), nrow=5, title=\"Random Cutout\")\nshow_img(aug_fn(example_monets), nrow=5, title=\"Differential Augmentation\")","metadata":{"executionInfo":{"elapsed":61,"status":"ok","timestamp":1685506351023,"user":{"displayName":"Garrett Devereux","userId":"06903598260790252274"},"user_tz":420},"id":"nDS_NO_qACoh","outputId":"733e9356-7fb6-400c-b28a-87e0f52625f4","execution":{"iopub.status.busy":"2023-06-01T00:55:44.588143Z","iopub.execute_input":"2023-06-01T00:55:44.589053Z","iopub.status.idle":"2023-06-01T00:55:47.526252Z","shell.execute_reply.started":"2023-06-01T00:55:44.589012Z","shell.execute_reply":"2023-06-01T00:55:47.525409Z"},"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","execution":{"iopub.status.busy":"2023-06-01T00:55:47.530511Z","iopub.execute_input":"2023-06-01T00:55:47.531714Z","iopub.status.idle":"2023-06-01T00:55:47.537768Z","shell.execute_reply.started":"2023-06-01T00:55:47.531667Z","shell.execute_reply":"2023-06-01T00:55:47.536894Z"},"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","execution":{"iopub.status.busy":"2023-06-01T00:55:47.539079Z","iopub.execute_input":"2023-06-01T00:55:47.539752Z","iopub.status.idle":"2023-06-01T00:55:47.550220Z","shell.execute_reply.started":"2023-06-01T00:55:47.539718Z","shell.execute_reply":"2023-06-01T00:55:47.549215Z"},"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":22,"status":"ok","timestamp":1685506351025,"user":{"displayName":"Garrett Devereux","userId":"06903598260790252274"},"user_tz":420},"id":"3AGFG97cB7yh","outputId":"44690386-341c-4ded-db75-51445520a35a","execution":{"iopub.status.busy":"2023-06-01T00:55:47.551685Z","iopub.execute_input":"2023-06-01T00:55:47.552277Z","iopub.status.idle":"2023-06-01T00:55:47.749779Z","shell.execute_reply.started":"2023-06-01T00:55:47.552243Z","shell.execute_reply":"2023-06-01T00:55:47.748770Z"},"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","execution":{"iopub.status.busy":"2023-06-01T00:55:47.751446Z","iopub.execute_input":"2023-06-01T00:55:47.751820Z","iopub.status.idle":"2023-06-01T00:55:47.760188Z","shell.execute_reply.started":"2023-06-01T00:55:47.751787Z","shell.execute_reply":"2023-06-01T00:55:47.758967Z"},"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":20,"status":"ok","timestamp":1685506351027,"user":{"displayName":"Garrett Devereux","userId":"06903598260790252274"},"user_tz":420},"id":"SECTyFKqHvG4","outputId":"46d44703-ce5d-4ee7-83cb-c2afb015bbdb","execution":{"iopub.status.busy":"2023-06-01T00:55:47.762287Z","iopub.execute_input":"2023-06-01T00:55:47.762698Z","iopub.status.idle":"2023-06-01T00:55:47.967128Z","shell.execute_reply.started":"2023-06-01T00:55:47.762661Z","shell.execute_reply":"2023-06-01T00:55:47.966100Z"},"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","execution":{"iopub.status.busy":"2023-06-01T00:55:47.969042Z","iopub.execute_input":"2023-06-01T00:55:47.969411Z","iopub.status.idle":"2023-06-01T00:55:47.983509Z","shell.execute_reply.started":"2023-06-01T00:55:47.969374Z","shell.execute_reply":"2023-06-01T00:55:47.982500Z"},"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":15641,"status":"ok","timestamp":1685506366652,"user":{"displayName":"Garrett Devereux","userId":"06903598260790252274"},"user_tz":420},"id":"7tZqDUCMKME8","outputId":"bb745565-a814-4350-94ff-ba785003dc4a","execution":{"iopub.status.busy":"2023-06-01T00:55:47.985050Z","iopub.execute_input":"2023-06-01T00:55:47.985555Z","iopub.status.idle":"2023-06-01T00:55:49.977189Z","shell.execute_reply.started":"2023-06-01T00:55:47.985520Z","shell.execute_reply":"2023-06-01T00:55:49.976256Z"},"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),\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","execution":{"iopub.status.busy":"2023-06-01T00:55:49.978849Z","iopub.execute_input":"2023-06-01T00:55:49.979199Z","iopub.status.idle":"2023-06-01T00:55:49.986599Z","shell.execute_reply.started":"2023-06-01T00:55:49.979166Z","shell.execute_reply":"2023-06-01T00:55:49.985539Z"},"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":68,"status":"ok","timestamp":1685506366653,"user":{"displayName":"Garrett Devereux","userId":"06903598260790252274"},"user_tz":420},"id":"b2XCBmpoUGdl","outputId":"4c40e716-b5ec-458e-d094-92f397929164","execution":{"iopub.status.busy":"2023-06-01T00:55:49.988194Z","iopub.execute_input":"2023-06-01T00:55:49.988920Z","iopub.status.idle":"2023-06-01T00:55:50.410830Z","shell.execute_reply.started":"2023-06-01T00:55:49.988887Z","shell.execute_reply":"2023-06-01T00:55:50.409654Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Dual Discriminators for Monets\n\nWe now will create a new Discriminator architecture inspired by Nguyen et. al in \"Dual Discriminator Generative Adversarial Nets\". Essentially, we will keep the first discrimintor architecture as our Photo Discriminator, but we will split the initial architecture up into two seperate parts for the Monet Discriminator - a body, and two heads. The body portion of the discriminator will be the exact same as the first portion above, minus the final convolution, mapping images from (batch_size, 3, 256, 256) -> (batch_size, 512, 31, 31). Then, the head completes the process, mapping the input from the body to a probability matrix of size  (batch_size, 1, 30, 30). \n\nThe main difference is that we apply two heads to the shared discriminator, each with different objectives (ex. BCE for Head 1 and Hinge loss for Head 2). This helps avoid the negative effects of overfitting the discriminator, as shared layers are disturbed with two different losses and discriminators are not able to overfit.\n\n\n![image3.jpg](https://raw.githubusercontent.com/garrettdevereux/DeepLearningGenerativeArt/main/D2Gan.png)\n\n[Image Source](https://arxiv.org/abs/1709.03831)","metadata":{"id":"MU8TuMGy-e-L"}},{"cell_type":"markdown","source":"## Body Discriminator\n\nThis acts as the shared discriminator for the two heads, and will map from(batch_size, 3, 256, 256) -> (batch_size, 512, 31, 31)","metadata":{"id":"fB2ORKVIwNG-"}},{"cell_type":"code","source":"class BodyDiscriminator(nn.Module):\n    def __init__(self, hid_channels, in_channels=3):\n        super().__init__()\n\n        # Initialize Series of Downsamples for the Body to map\n        # (batch_size, 3, 256, 256) -> (batch_size, 512, 31, 31)\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        )\n\n        # Apply Init Weights\n        self.block.apply(init_weights)\n        \n    def forward(self, x):\n        return self.block(x)","metadata":{"id":"q24vXDFE-k34","execution":{"iopub.status.busy":"2023-06-01T00:55:50.412641Z","iopub.execute_input":"2023-06-01T00:55:50.413005Z","iopub.status.idle":"2023-06-01T00:55:50.420582Z","shell.execute_reply.started":"2023-06-01T00:55:50.412972Z","shell.execute_reply":"2023-06-01T00:55:50.419537Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Quick Sanity Check to Ensure the Body discriminator operates on our image size correctly, generating an intermediate output of size 512 x 31 x 31 for the Heads.","metadata":{"id":"XeG4zvp3wqAN"}},{"cell_type":"code","source":"def test_Body():\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 = BodyDiscriminator(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(\"Intermediate Body Shape:\")\n  print(probability_scores.size())\n  print()\n\ntest_Body()","metadata":{"executionInfo":{"elapsed":63,"status":"ok","timestamp":1685506366654,"user":{"displayName":"Garrett Devereux","userId":"06903598260790252274"},"user_tz":420},"id":"l-y3yEg2AdPd","outputId":"af7b4073-6bb5-4667-e4a3-2923f558dadb","execution":{"iopub.status.busy":"2023-06-01T00:55:50.422518Z","iopub.execute_input":"2023-06-01T00:55:50.423444Z","iopub.status.idle":"2023-06-01T00:55:50.819861Z","shell.execute_reply.started":"2023-06-01T00:55:50.423409Z","shell.execute_reply":"2023-06-01T00:55:50.818767Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Head Discriminator\n\nThis acts as the Head discriminator that will be initialized twice at the end of the Monet Body Discriminator, and will map from(batch_size, 512, 31, 31) -> to our probability matrix of size (batch_size, 1, 30, 30)","metadata":{"id":"PLy7-PpNw8iO"}},{"cell_type":"code","source":"class HeadDiscriminator(nn.Module):\n    def __init__(self, hid_channels, in_channels=3):\n        super().__init__()\n\n        # Final Up Convolution to map to probability matrix\n        # (batch_size, 512, 31, 31) -> (batch_size, 1, 30, 30)\n        self.block = nn.Sequential(\n            # Final Convolution mapping\n            # (512, 31, 31) -> (1, 30, 30)\n            nn.Conv2d(hid_channels*8, 1, kernel_size=4, padding=1),\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":"szdHoWGXAGyb","execution":{"iopub.status.busy":"2023-06-01T00:55:50.821442Z","iopub.execute_input":"2023-06-01T00:55:50.821843Z","iopub.status.idle":"2023-06-01T00:55:50.828720Z","shell.execute_reply.started":"2023-06-01T00:55:50.821807Z","shell.execute_reply":"2023-06-01T00:55:50.827505Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Quick Sanity Check to Ensure the Head discriminator operates on our image size correctly, generating a final probability matrix of size 1 x 30 x 30 for the Heads.","metadata":{"id":"_UlbxoPUxfhP"}},{"cell_type":"code","source":"def test_Head():\n  # Test with input size 512 x 31 x 31 and batch 5\n  in_channels = 512\n  hidden_channels = 64\n  out_channels = 3\n  batch_size = 5\n  H = 31\n  W = 31\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 = HeadDiscriminator(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\ntest_Head()","metadata":{"executionInfo":{"elapsed":59,"status":"ok","timestamp":1685506366655,"user":{"displayName":"Garrett Devereux","userId":"06903598260790252274"},"user_tz":420},"id":"L8rIDVJdA-9u","outputId":"c8b28c0e-9980-4b80-a2e6-55c3e8f6259a","execution":{"iopub.status.busy":"2023-06-01T00:55:50.830304Z","iopub.execute_input":"2023-06-01T00:55:50.831340Z","iopub.status.idle":"2023-06-01T00:55:50.855643Z","shell.execute_reply.started":"2023-06-01T00:55:50.831258Z","shell.execute_reply":"2023-06-01T00:55:50.854613Z"},"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":55,"status":"ok","timestamp":1685506366655,"user":{"displayName":"Garrett Devereux","userId":"06903598260790252274"},"user_tz":420},"id":"3LvRnlEVeaPW","outputId":"12ec7434-8aa5-4bbf-934f-07dffacb0cc5","execution":{"iopub.status.busy":"2023-06-01T00:55:50.857087Z","iopub.execute_input":"2023-06-01T00:55:50.857510Z","iopub.status.idle":"2023-06-01T00:55:56.579115Z","shell.execute_reply.started":"2023-06-01T00:55:50.857462Z","shell.execute_reply":"2023-06-01T00:55:56.576547Z"},"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 photos\n# Initial PatchGan architecture\nphoto_discriminator = Discriminator(hid_channels)\n\n\n# TWO OBJECTIVE MONET DISCRIMINATOR\nmonet_discriminator = BodyDiscriminator(hid_channels)\nmonet_dhead1 = HeadDiscriminator(hid_channels)\nmonet_dhead2 = HeadDiscriminator(hid_channels)\n\n# Give a real photo to each discriminator (first one photo shown above)\nreal_photo_scores = photo_discriminator(example_photos[0])\n\n# Note we will get two results and sent to each head\nintermediate = monet_discriminator(example_monets[0])\nreal_monet_scores1 = monet_dhead1(intermediate)\nreal_monet_scores2 = monet_dhead2(intermediate)\n\n# Sizes are randomly initialized\nprint(\"Shapes\")\nprint(\"Photo Discriminator Size\")\nprint(real_photo_scores.shape)\nprint()\nprint(\"Monet Discriminator Size\")\nprint(\"Intermediate\")\nprint(intermediate.shape)\nprint(\"After Heads\")\nprint(real_monet_scores1.shape)\nprint(real_monet_scores2.shape)\nprint()\nprint(\"Actual Scores\")\nprint(real_photo_scores)\nprint(real_monet_scores1)\nprint(real_monet_scores2)","metadata":{"executionInfo":{"elapsed":31,"status":"ok","timestamp":1685506366655,"user":{"displayName":"Garrett Devereux","userId":"06903598260790252274"},"user_tz":420},"id":"EqXpm9eaf0UY","outputId":"da72b326-a039-42be-8577-942adae18a82","execution":{"iopub.status.busy":"2023-06-01T00:55:56.582863Z","iopub.execute_input":"2023-06-01T00:55:56.587750Z","iopub.status.idle":"2023-06-01T00:55:56.965406Z","shell.execute_reply.started":"2023-06-01T00:55:56.587708Z","shell.execute_reply":"2023-06-01T00:55:56.962999Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Build THE CycleGAN!\n\nWe now have our Generator and Discriminators 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. With a new architecture, we will have two loss functions for each generator and discriminator corresponding to the two Heads: one for Hinge, and one for BCE. Cycle and Identity Loss remian the same as the baseline.","metadata":{"id":"yAw5sa_IoAVy"}},{"cell_type":"markdown","source":"### Generator Loss\n\nOur Generator Losses for in reference to the two discriminator Heads","metadata":{"id":"PPjVeQXAhlim"}},{"cell_type":"code","source":"# Gen Loss 1 for Hinge\ndef generator_loss1(fake_discriminator_score):\n    real_labels = torch.ones_like(fake_discriminator_score)\n    # Retrun mse between actual scores and 1s\n    return F.mse_loss(fake_discriminator_score, real_labels)\n\n# Gen Loss 2 using BCE Loss\ndef generator_loss2(fake_discriminator_score):\n    return torch.mean(torch.nn.BCEWithLogitsLoss()(fake_discriminator_score, torch.ones_like(fake_discriminator_score)))","metadata":{"id":"YLh8EtJOktSQ","execution":{"iopub.status.busy":"2023-06-01T00:55:56.966968Z","iopub.execute_input":"2023-06-01T00:55:56.967331Z","iopub.status.idle":"2023-06-01T00:55:56.973639Z","shell.execute_reply.started":"2023-06-01T00:55:56.967298Z","shell.execute_reply":"2023-06-01T00:55:56.972397Z"},"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","execution":{"iopub.status.busy":"2023-06-01T00:55:56.975883Z","iopub.execute_input":"2023-06-01T00:55:56.976296Z","iopub.status.idle":"2023-06-01T00:55:56.985089Z","shell.execute_reply.started":"2023-06-01T00:55:56.976212Z","shell.execute_reply":"2023-06-01T00:55:56.984119Z"},"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_idt):\n    # Simply return reg * L1 loss between real and identity\n    return lambda_idt * 0.5 * F.l1_loss(identity_image, real)","metadata":{"id":"PBxx6ehUmlYH","execution":{"iopub.status.busy":"2023-06-01T00:55:56.986882Z","iopub.execute_input":"2023-06-01T00:55:56.987331Z","iopub.status.idle":"2023-06-01T00:55:56.993810Z","shell.execute_reply.started":"2023-06-01T00:55:56.987298Z","shell.execute_reply":"2023-06-01T00:55:56.992883Z"},"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":"markdown","source":"### Discriminator Loss\n\nOur Discriminator loss computes 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. As we our now doing Dual Objective Discriminators, we have corresponding loss functions for both Hinge Loss and BCE.\n\n![image7.png](https://raw.githubusercontent.com/garrettdevereux/DeepLearningGenerativeArt/main/DiscriminatorLoss.png)","metadata":{"id":"z0HIvxSzoCz9"}},{"cell_type":"code","source":"# MSE Discriminator Loss\ndef discriminator_loss1(real_discriminator_score, fake_discriminator_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_discriminator_score)\n    # Return mse between actual scores and 1s\n    real_loss = F.mse_loss(real_discriminator_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_discriminator_score)\n    # Return mse between actual scores and 1s\n    fake_loss = F.mse_loss(fake_discriminator_score, fake_labels)\n        \n    # Return average of losses\n    disc_loss = (fake_loss + real_loss) * 0.5\n    return disc_loss\n\n# BCE Discriminator Loss\ndef discriminator_loss2(real_discriminator_score, fake_discriminator_score):\n    generated_loss = torch.nn.BCEWithLogitsLoss()(fake_discriminator_score, torch.zeros_like(fake_discriminator_score))\n    real_loss = torch.nn.BCEWithLogitsLoss()(real_discriminator_score, torch.ones_like(real_discriminator_score))\n    total_disc_loss = real_loss + generated_loss\n\n    return torch.mean(total_disc_loss * 0.5)","metadata":{"executionInfo":{"elapsed":2,"status":"ok","timestamp":1685506703218,"user":{"displayName":"Garrett Devereux","userId":"06903598260790252274"},"user_tz":420},"id":"CZMXvaBmn9Sg","execution":{"iopub.status.busy":"2023-06-01T00:55:57.000099Z","iopub.execute_input":"2023-06-01T00:55:57.000380Z","iopub.status.idle":"2023-06-01T00:55:57.007287Z","shell.execute_reply.started":"2023-06-01T00:55:57.000357Z","shell.execute_reply":"2023-06-01T00:55:57.006259Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"For Photo Discriminator","metadata":{"id":"11rXBWPh9xbn"}},{"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        decay_epochs,\n        monet_gen = None,\n        photo_gen = None,\n        monet_disc = None,\n        photo_disc = None, \n        monet_disc_head1 = None,\n        monet_disc_head2 = 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\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        # DUAL OBJECTIVE DISCRIMINATORS\n        if monet_disc == None:\n          self.monet_discriminator = BodyDiscriminator(hid_channels)\n        else:\n          self.monet_discriminator = monet_disc\n\n        # HEADS FOR MONET DISCRIM\n\n        if monet_disc_head1 == None:\n          self.monet_disc_head1 = HeadDiscriminator(hid_channels)\n        else:\n          self.monet_disc_head1 = HeadDiscriminator(hid_channels)\n\n        if monet_disc_head2 == None:\n          self.monet_disc_head2 = HeadDiscriminator(hid_channels)\n        else:\n          self.monet_disc_head2 = HeadDiscriminator(hid_channels)\n\n        # Metrics for Loss graphs\n        self.losses = {\n            \"gen_loss\": [],\n            \"disc_loss_M\": [],\n            \"disc_loss_P\": [],\n        }\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    # 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\n        # Set up optimizer for discriminators\n        opt_disc = self.optimizer(\n            list(self.photo_discriminator.parameters()) + list(self.monet_discriminator.parameters()) + list(self.monet_disc_head1.parameters()) + list(self.monet_disc_head2.parameters()),\n            **opt_config,\n        )\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 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        batch_size = float(real_monet.shape[0])\n\n        # Store the optimizers for the generators and discriminators\n        gen_optimizer, disc_optimizer = self.optimizers()\n\n        ## PHOTO CONSTRUCTION: fakes, reconstructions, and identities\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\n        ## DIFFERENTIAL AUGMENTATION!!!!!!\n        \n        # Pass a batch of real and fake to diff aug funciton\n        aug_real_monet = aug_fn(real_monet)\n        aug_fake_monet = aug_fn(fake_monet)\n\n        ## DISCRIMINATOR SCORES\n        # TWO OBJECTIVE DISCRIMINATOR!!!!!!!!!\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\n        # HEAD 1 of MONET DISCRIMINATOR\n        # Score of augmented fake from monet discriminator head 1\n        fake_monet_discrim_score1 = self.monet_disc_head1(self.monet_discriminator(aug_fake_monet))\n        # Score of augmented real from monet discriminator head 1\n        real_monet_discrim_score1 = self.monet_disc_head1(self.monet_discriminator(aug_real_monet))\n\n        # HEAD 2 of MONET DISCRIMINATOR\n        # Score of augmented fake from monet discriminator head 2\n        fake_monet_discrim_score2 = self.monet_disc_head1(self.monet_discriminator(aug_fake_monet))\n        # Score of augmented real from monet discriminator head 2\n        real_monet_discrim_score2 = self.monet_disc_head1(self.monet_discriminator(aug_real_monet))\n\n        # PHOTO DISCRIMINATOR SCORES\n        # Scores of how good the real photos are\n        real_photo_discrim_score = self.photo_discriminator(real_photo)\n        # Scores of how good the fake generated photos are\n        fake_photo_discrim_score = self.photo_discriminator(fake_photo)\n\n\n        ## COMPUTE LOSSES AND OPTIMIZE\n\n        # Standard pipeline: Loss -> zero_grad -> backward -> step\n        \n        if self.step_num % 2 == 0:\n\n            # GENERATOR LOSS\n            self.toggle_optimizer(gen_optimizer)  \n\n            # GENERATOR LOSS\n            monet_base_gen_loss = 0.4 * (generator_loss1(fake_monet_discrim_score1) + generator_loss2(fake_monet_discrim_score2))\n\n            photo_base_gen_loss = generator_loss1(fake_photo_discrim_score)\n\n            total_cycle_loss = get_cycle_loss(real_monet, reconstructed_monet, (self.hparams.lambda_cycle[0] / batch_size)) + get_cycle_loss(real_photo, reconstructed_photo, (self.hparams.lambda_cycle[1] / batch_size))\n\n            total_monet_gen_loss = monet_base_gen_loss + total_cycle_loss + get_idt_loss(real_monet, monet_identity, self.hparams.lambda_idt / batch_size)\n            total_photo_gen_loss = photo_base_gen_loss + total_cycle_loss + get_idt_loss(real_photo, photo_identity, self.hparams.lambda_idt / batch_size)\n\n            total_generator_loss = total_monet_gen_loss + total_photo_gen_loss\n            # Complete Pipeline\n            gen_optimizer.zero_grad()\n            self.manual_backward(total_generator_loss)\n            gen_optimizer.step()\n            self.untoggle_optimizer(gen_optimizer)\n\n\n            # DISCRIMINATOR LOSS\n            self.toggle_optimizer(disc_optimizer)\n\n            # MONET DISCRIMINATOR LOSS\n            monet_head_loss1 = discriminator_loss1(real_monet_discrim_score1, fake_monet_discrim_score1.detach())\n            monet_head_loss2 = discriminator_loss2(real_monet_discrim_score2, fake_monet_discrim_score2.detach())\n            monet_disc_loss = monet_head_loss1 + monet_head_loss2\n\n            # PHOTO DISCRIMINATOR LOSS\n            photo_disc_loss = discriminator_loss1(real_photo_discrim_score, self.photo_discriminator(fake_photo.detach()))\n\n            # Complete Pipeline\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(disc_optimizer)\n            \n        else:\n            # DISCRIMINATOR LOSS\n            self.toggle_optimizer(disc_optimizer)\n\n            # MONET DISCRIMINATOR LOSS\n            monet_head_loss1 = discriminator_loss1(real_monet_discrim_score1, fake_monet_discrim_score1)\n            monet_head_loss2 = discriminator_loss2(real_monet_discrim_score2, fake_monet_discrim_score2)\n            monet_disc_loss = monet_head_loss1 + monet_head_loss2\n\n            # PHOTO DISCRIMINATOR LOSS\n            photo_disc_loss = discriminator_loss1(real_photo_discrim_score, fake_photo_discrim_score)\n\n            # Complete Pipeline\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(disc_optimizer)\n            \n            # GENERATOR LOSS\n            self.toggle_optimizer(gen_optimizer)  \n\n            # GENERATOR LOSS\n            monet_base_gen_loss = 0.4 * (generator_loss1(fake_monet_discrim_score1.detach()) + generator_loss2(fake_monet_discrim_score2.detach()))\n\n            photo_base_gen_loss = generator_loss1(self.photo_discriminator(fake_photo.detach()))\n\n            total_cycle_loss = get_cycle_loss(real_monet, reconstructed_monet, (self.hparams.lambda_cycle[0] / batch_size)) + get_cycle_loss(real_photo, reconstructed_photo, (self.hparams.lambda_cycle[1] / batch_size))\n\n            total_monet_gen_loss = monet_base_gen_loss + total_cycle_loss + get_idt_loss(real_monet, monet_identity, self.hparams.lambda_idt / batch_size)\n            total_photo_gen_loss = photo_base_gen_loss + total_cycle_loss + get_idt_loss(real_photo, photo_identity, self.hparams.lambda_idt / batch_size)\n\n            total_generator_loss = total_monet_gen_loss + total_photo_gen_loss\n            # Complete Pipeline\n            gen_optimizer.zero_grad()\n            self.manual_backward(total_generator_loss)\n            gen_optimizer.step()\n            self.untoggle_optimizer(gen_optimizer)\n            \n\n        \n        # Done with training step\n        # Log metrics and losses\n        metrics = {\n            \"gen_loss\": total_generator_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        self.step += 1\n\n        # Add to Losses for graphs\n        self.losses[\"gen_loss\"].append(float(total_generator_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        \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.step = 0\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        for sch in self.lr_schedulers():\n          sch.step()\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\n        dhead1_path = checkpoints + \"/dhead1_discriminator\" + str(self.current_epoch) + \".pth\"\n        torch.save(self.monet_disc_head1.state_dict(), dhead1_path)\n\n        dhead2_path = checkpoints + \"/dhead2_discriminator\" + str(self.current_epoch) + \".pth\"\n        torch.save(self.monet_disc_head2.state_dict(), dhead2_path)\n        print(\"Training ended.\")\n        \"\"\"\n        \n    def on_predict_epoch_end(self):\n        predictions = self.trainer.predict_loop.predictions\n        print(f\"Done predicting\")","metadata":{"executionInfo":{"elapsed":453,"status":"ok","timestamp":1685506706539,"user":{"displayName":"Garrett Devereux","userId":"06903598260790252274"},"user_tz":420},"id":"VIZx1_P6hreW","execution":{"iopub.status.busy":"2023-06-01T00:55:57.009188Z","iopub.execute_input":"2023-06-01T00:55:57.009994Z","iopub.status.idle":"2023-06-01T00:55:57.053711Z","shell.execute_reply.started":"2023-06-01T00:55:57.009959Z","shell.execute_reply":"2023-06-01T00:55:57.052719Z"},"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":"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\": 16\n}","metadata":{"executionInfo":{"elapsed":380,"status":"ok","timestamp":1685506715203,"user":{"displayName":"Garrett Devereux","userId":"06903598260790252274"},"user_tz":420},"id":"42viknna6hhw","execution":{"iopub.status.busy":"2023-06-01T01:10:44.502926Z","iopub.execute_input":"2023-06-01T01:10:44.503303Z","iopub.status.idle":"2023-06-01T01:10:44.509526Z","shell.execute_reply.started":"2023-06-01T01:10:44.503271Z","shell.execute_reply":"2023-06-01T01:10:44.508587Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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\": 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 number of epochs for training\n    \"num_epochs\": 25,\n\n    \"decay_epochs\": 20,\n}","metadata":{"executionInfo":{"elapsed":11,"status":"ok","timestamp":1685506797968,"user":{"displayName":"Garrett Devereux","userId":"06903598260790252274"},"user_tz":420},"id":"QT_DFMILg-aS","execution":{"iopub.status.busy":"2023-06-01T01:10:38.900770Z","iopub.execute_input":"2023-06-01T01:10:38.901155Z","iopub.status.idle":"2023-06-01T01:10:38.908344Z","shell.execute_reply.started":"2023-06-01T01:10:38.901125Z","shell.execute_reply":"2023-06-01T01:10:38.907189Z"},"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\": 1,\n}","metadata":{"executionInfo":{"elapsed":409,"status":"ok","timestamp":1685506746802,"user":{"displayName":"Garrett Devereux","userId":"06903598260790252274"},"user_tz":420},"id":"JnVIXbjtjmTQ","execution":{"iopub.status.busy":"2023-06-01T00:55:57.078954Z","iopub.execute_input":"2023-06-01T00:55:57.079341Z","iopub.status.idle":"2023-06-01T00:55:57.088522Z","shell.execute_reply.started":"2023-06-01T00:55:57.079310Z","shell.execute_reply":"2023-06-01T00:55:57.087433Z"},"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":"4f9ff6f9-51e1-4aeb-c275-7aec3c5ac8fc","execution":{"iopub.status.busy":"2023-06-01T00:57:32.874328Z","iopub.execute_input":"2023-06-01T00:57:32.875276Z","iopub.status.idle":"2023-06-01T01:06:35.529983Z","shell.execute_reply.started":"2023-06-01T00:57:32.875239Z","shell.execute_reply":"2023-06-01T01:06:35.528842Z"},"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\"]))\n\n# Plot the Generator Loss\naxes[0].set_title(\"Generator Loss\")\naxes[0].plot(training_steps, model.losses[\"gen_loss\"], label='Monet Generator Loss')\naxes[0].legend()\naxes[0].set_xlabel('Steps')\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('Steps')\n\nplt.show()","metadata":{"executionInfo":{"elapsed":1408,"status":"ok","timestamp":1685505261152,"user":{"displayName":"Garrett Devereux","userId":"06903598260790252274"},"user_tz":420},"id":"VJE4FwVGF36B","outputId":"90355949-3319-446b-becd-9c7778eb8efc","execution":{"iopub.status.busy":"2023-06-01T00:57:18.722344Z","iopub.execute_input":"2023-06-01T00:57:18.722739Z","iopub.status.idle":"2023-06-01T00:57:20.028128Z","shell.execute_reply.started":"2023-06-01T00:57:18.722691Z","shell.execute_reply":"2023-06-01T00:57:20.026986Z"},"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\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\n\nreal_monet_score1 = model.monet_disc_head1(model.monet_discriminator(example_monets[0]))\nreal_monet_score2 = model.monet_disc_head1(model.monet_discriminator(example_monets[0]))\nreal_photo_score = model.photo_discriminator(example_photos[0])\nfake_monet_score = model.monet_disc_head1(model.monet_discriminator(fake_monets[0]))\nfake_photo_score = model.photo_discriminator(fake_photos[0])\nones_label = torch.ones_like(real_photo_score)\nzeros_label = torch.zeros_like(real_photo_score)\n\nprint(\"Visualizing Results produced from Original Photos\")\n# Visualize from Original Photos\nshow_img(example_photos, nrow=5, title=\"Original Photos\")\nshow_img(fake_monets, nrow=5, title=\"Generated Monets\")\nshow_img(recon_photo, nrow=5, title=\"Reconstructed Photos\")\nshow_img(idt_photo, nrow=5, title=\"Identity Photos\")\n\n# Visualize from Original Monets\nprint(\"Visualizing Results produced from Original Monets\")\nshow_img(example_monets, nrow=5, title=\"Original Monets\")\nshow_img(fake_photos, nrow=5, title=\"Generated Photos\")\nshow_img(recon_monet, nrow=5, title=\"Reconstructed Monets\")\nshow_img(idt_monet, nrow=5, title=\"Identity Monets\")\n\nprint(\"Discrim Scores\")\nprint(real_monet_score1)\nprint(real_monet_score2)\nprint(real_photo_score)","metadata":{"executionInfo":{"elapsed":15423,"status":"ok","timestamp":1685430323454,"user":{"displayName":"Garrett Devereux","userId":"06903598260790252274"},"user_tz":420},"id":"y0OqwP7xrNfA","outputId":"4ace6250-7d92-4916-a85b-83c3c9b992b9","execution":{"iopub.status.busy":"2023-06-01T01:07:08.611816Z","iopub.execute_input":"2023-06-01T01:07:08.612202Z","iopub.status.idle":"2023-06-01T01:07:20.233382Z","shell.execute_reply.started":"2023-06-01T01:07:08.612162Z","shell.execute_reply":"2023-06-01T01:07:20.232336Z"},"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_understand)","metadata":{"execution":{"iopub.status.busy":"2023-06-01T01:07:32.034754Z","iopub.execute_input":"2023-06-01T01:07:32.035130Z","iopub.status.idle":"2023-06-01T01:08:56.089316Z","shell.execute_reply.started":"2023-06-01T01:07:32.035097Z","shell.execute_reply":"2023-06-01T01:08:56.088191Z"},"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":{"execution":{"iopub.status.busy":"2023-06-01T01:09:32.144840Z","iopub.execute_input":"2023-06-01T01:09:32.145208Z","iopub.status.idle":"2023-06-01T01:09:53.803252Z","shell.execute_reply.started":"2023-06-01T01:09:32.145176Z","shell.execute_reply":"2023-06-01T01:09:53.801499Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"checkpoints = \"/kaggle/working/\"\n\nmonet_gen_path = checkpoints + \"/monet_generator\" + \".pth\"\ntorch.save(model.monet_generator.state_dict(), monet_gen_path)\n\nphoto_gen_path = checkpoints + \"/photo_generator\" + \".pth\"\ntorch.save(model.photo_generator.state_dict(), photo_gen_path)\n\nmonet_disc_path = checkpoints + \"/monet_discriminator\" + \".pth\"\ntorch.save(model.monet_discriminator.state_dict(), monet_disc_path)\n\nphoto_disc_path = checkpoints + \"/photo_discriminator\" + \".pth\"\ntorch.save(model.photo_discriminator.state_dict(), photo_disc_path)\n\ndhead1_path = checkpoints + \"/dhead1_discriminator\" + \".pth\"\ntorch.save(model.monet_disc_head1.state_dict(), dhead1_path)\ndhead2_path = checkpoints + \"/dhead2_discriminator\" + \".pth\"\ntorch.save(model.monet_disc_head2.state_dict(), dhead2_path)","metadata":{"execution":{"iopub.status.busy":"2023-06-01T01:09:57.593582Z","iopub.execute_input":"2023-06-01T01:09:57.593944Z","iopub.status.idle":"2023-06-01T01:09:58.114838Z","shell.execute_reply.started":"2023-06-01T01:09:57.593914Z","shell.execute_reply":"2023-06-01T01:09:58.113767Z"},"trusted":true},"execution_count":null,"outputs":[]}]}