{"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":"code","source":"!pip install pytorch_lightning\n!pip install deepspeed","metadata":{"executionInfo":{"elapsed":21727,"status":"ok","timestamp":1685393934976,"user":{"displayName":"Begona Garcia Malaxechebarria","userId":"06396578984430635223"},"user_tz":420},"id":"-mQnQb6rm1C7","outputId":"6b1d8e66-4250-4a8c-cb48-d64ca1f01f0b","execution":{"iopub.status.busy":"2023-06-15T09:14:22.163927Z","iopub.execute_input":"2023-06-15T09:14:22.164742Z","iopub.status.idle":"2023-06-15T09:14:44.747021Z","shell.execute_reply.started":"2023-06-15T09:14:22.164699Z","shell.execute_reply":"2023-06-15T09:14:44.745831Z"},"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\n\ndevice = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\nprint(device)","metadata":{"executionInfo":{"elapsed":5204,"status":"ok","timestamp":1685393940158,"user":{"displayName":"Begona Garcia Malaxechebarria","userId":"06396578984430635223"},"user_tz":420},"id":"zXMDumZap7i4","outputId":"ae2b8b2d-f8f0-4329-b706-45e5c3019bce","execution":{"iopub.status.busy":"2023-06-15T09:14:44.749717Z","iopub.execute_input":"2023-06-15T09:14:44.75129Z","iopub.status.idle":"2023-06-15T09:14:44.759751Z","shell.execute_reply.started":"2023-06-15T09:14:44.751252Z","shell.execute_reply":"2023-06-15T09:14:44.75884Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CustomImageDataset(Dataset):\n    def __init__(self, filenames, stage):\n        self.filenames = filenames\n        self.stage = stage\n        \n    def __len__(self):\n        return len(self.filenames)\n    \n    def __getitem__(self, idx):\n        img_name = self.filenames[idx]\n        img = (read_image(img_name) / 255.0) * 2 - 1\n        return img","metadata":{"id":"eZGw_PPTKLh2","executionInfo":{"status":"ok","timestamp":1685393976135,"user_tz":420,"elapsed":24,"user":{"displayName":"Begona Garcia Malaxechebarria","userId":"06396578984430635223"}},"execution":{"iopub.status.busy":"2023-06-15T09:14:44.761138Z","iopub.execute_input":"2023-06-15T09:14:44.761459Z","iopub.status.idle":"2023-06-15T09:14:44.775437Z","shell.execute_reply.started":"2023-06-15T09:14:44.76143Z","shell.execute_reply":"2023-06-15T09:14:44.774535Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CustomDataModule(L.LightningDataModule):\n    def __init__(\n        self,\n        monet_dir,\n        photo_dir,\n        loader_config,\n        sample_size,\n        batch_size,\n    ):\n        super().__init__()\n        self.loader_config = loader_config\n        self.sample_size = sample_size\n        self.batch_size = batch_size\n\n        self.monet_filenames = sorted(glob.glob(monet_dir))\n        self.photo_filenames = sorted(glob.glob(photo_dir))\n        \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            self.valid_photo = CustomImageDataset(self.photo_filenames, None)\n\n    def train_dataloader(self):\n        loader_config = {\n            \"shuffle\": True,\n            \"drop_last\": True,\n            \"batch_size\": self.batch_size,\n            **self.loader_config,\n        }\n        loaders = {\n            \"monet\": DataLoader(self.train_monet, **loader_config),\n            \"photo\": DataLoader(self.train_photo, **loader_config)\n        }\n        return loaders\n    \n    def val_dataloader(self):\n        return DataLoader(\n            self.valid_photo,\n            batch_size=self.sample_size,\n            **self.loader_config,\n        )\n    \n    def test_dataloader(self):\n        return self.val_dataloader()\n    \n    def predict_dataloader(self):\n        return DataLoader(\n            self.valid_photo,\n            batch_size=self.batch_size,\n            **self.loader_config,\n        )\n","metadata":{"id":"Yno-7MPvMOhI","executionInfo":{"status":"ok","timestamp":1685393976136,"user_tz":420,"elapsed":20,"user":{"displayName":"Begona Garcia Malaxechebarria","userId":"06396578984430635223"}},"execution":{"iopub.status.busy":"2023-06-15T09:14:44.779818Z","iopub.execute_input":"2023-06-15T09:14:44.780112Z","iopub.status.idle":"2023-06-15T09:14:44.792413Z","shell.execute_reply.started":"2023-06-15T09:14:44.780089Z","shell.execute_reply":"2023-06-15T09:14:44.791533Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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","executionInfo":{"status":"ok","timestamp":1685393976138,"user_tz":420,"elapsed":19,"user":{"displayName":"Begona Garcia Malaxechebarria","userId":"06396578984430635223"}},"execution":{"iopub.status.busy":"2023-06-15T09:14:44.793852Z","iopub.execute_input":"2023-06-15T09:14:44.794187Z","iopub.status.idle":"2023-06-15T09:14:44.805315Z","shell.execute_reply.started":"2023-06-15T09:14:44.794157Z","shell.execute_reply":"2023-06-15T09:14:44.804369Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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":47794,"status":"ok","timestamp":1685394023914,"user":{"displayName":"Begona Garcia Malaxechebarria","userId":"06396578984430635223"},"user_tz":420},"id":"GnmReHHjQsZF","outputId":"260f9083-8312-4ab1-9ccc-3fe2b675657d","execution":{"iopub.status.busy":"2023-06-15T09:14:44.807621Z","iopub.execute_input":"2023-06-15T09:14:44.80796Z","iopub.status.idle":"2023-06-15T09:14:46.658097Z","shell.execute_reply.started":"2023-06-15T09:14:44.807931Z","shell.execute_reply":"2023-06-15T09:14:46.656797Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def 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\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\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\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\nAUGMENT_FNS = {\n    'color': [rand_brightness, rand_saturation, rand_contrast],\n}\n\ndef aug_fn(image):\n    return DiffAugment(image, \"color\")\n","metadata":{"execution":{"iopub.status.busy":"2023-06-15T09:14:46.659991Z","iopub.execute_input":"2023-06-15T09:14:46.660661Z","iopub.status.idle":"2023-06-15T09:14:46.673308Z","shell.execute_reply.started":"2023-06-15T09:14:46.660619Z","shell.execute_reply":"2023-06-15T09:14:46.671543Z"},"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()","metadata":{"executionInfo":{"elapsed":10021,"status":"ok","timestamp":1685394033854,"user":{"displayName":"Begona Garcia Malaxechebarria","userId":"06396578984430635223"},"user_tz":420},"id":"SecQv3LeRnb-","outputId":"247b7126-dc72-45fc-ded3-d6da7bd13eb9","execution":{"iopub.status.busy":"2023-06-15T09:14:46.675089Z","iopub.execute_input":"2023-06-15T09:14:46.675694Z","iopub.status.idle":"2023-06-15T09:14:46.686079Z","shell.execute_reply.started":"2023-06-15T09:14:46.67566Z","shell.execute_reply":"2023-06-15T09:14:46.685064Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"show_img(example_monets, nrow=5, title=\"Monet Tabloları\")\nshow_img(example_photos, nrow=5, title=\"Fotoğraflar\")","metadata":{"execution":{"iopub.status.busy":"2023-06-15T09:14:46.687466Z","iopub.execute_input":"2023-06-15T09:14:46.688305Z","iopub.status.idle":"2023-06-15T09:14:47.645953Z","shell.execute_reply.started":"2023-06-15T09:14:46.688272Z","shell.execute_reply":"2023-06-15T09:14:47.644683Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Visualize with Differential Augmentation","metadata":{}},{"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(aug_fn(example_monets), nrow=5, title=\"Differential Augmentation\")","metadata":{"execution":{"iopub.status.busy":"2023-06-15T09:14:47.650529Z","iopub.execute_input":"2023-06-15T09:14:47.651221Z","iopub.status.idle":"2023-06-15T09:14:49.64015Z","shell.execute_reply.started":"2023-06-15T09:14:47.651188Z","shell.execute_reply":"2023-06-15T09:14:49.639271Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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","executionInfo":{"status":"ok","timestamp":1685394033857,"user_tz":420,"elapsed":62,"user":{"displayName":"Begona Garcia Malaxechebarria","userId":"06396578984430635223"}},"execution":{"iopub.status.busy":"2023-06-15T09:14:49.64175Z","iopub.execute_input":"2023-06-15T09:14:49.642409Z","iopub.status.idle":"2023-06-15T09:14:49.648675Z","shell.execute_reply.started":"2023-06-15T09:14:49.642376Z","shell.execute_reply":"2023-06-15T09:14:49.647778Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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","executionInfo":{"status":"ok","timestamp":1685394033859,"user_tz":420,"elapsed":59,"user":{"displayName":"Begona Garcia Malaxechebarria","userId":"06396578984430635223"}},"execution":{"iopub.status.busy":"2023-06-15T09:14:49.650172Z","iopub.execute_input":"2023-06-15T09:14:49.650761Z","iopub.status.idle":"2023-06-15T09:14:49.660165Z","shell.execute_reply.started":"2023-06-15T09:14:49.65073Z","shell.execute_reply":"2023-06-15T09:14:49.659391Z"},"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":60,"status":"ok","timestamp":1685394033862,"user":{"displayName":"Begona Garcia Malaxechebarria","userId":"06396578984430635223"},"user_tz":420},"id":"3AGFG97cB7yh","outputId":"2e50debb-26f3-4589-b6b6-e9fc84eb82b2","execution":{"iopub.status.busy":"2023-06-15T09:14:49.661619Z","iopub.execute_input":"2023-06-15T09:14:49.662322Z","iopub.status.idle":"2023-06-15T09:14:49.808259Z","shell.execute_reply.started":"2023-06-15T09:14:49.662291Z","shell.execute_reply":"2023-06-15T09:14:49.807259Z"},"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","executionInfo":{"status":"ok","timestamp":1685394033864,"user_tz":420,"elapsed":45,"user":{"displayName":"Begona Garcia Malaxechebarria","userId":"06396578984430635223"}},"execution":{"iopub.status.busy":"2023-06-15T09:14:49.809974Z","iopub.execute_input":"2023-06-15T09:14:49.810621Z","iopub.status.idle":"2023-06-15T09:14:49.819684Z","shell.execute_reply.started":"2023-06-15T09:14:49.810584Z","shell.execute_reply":"2023-06-15T09:14:49.818796Z"},"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":45,"status":"ok","timestamp":1685394033866,"user":{"displayName":"Begona Garcia Malaxechebarria","userId":"06396578984430635223"},"user_tz":420},"id":"SECTyFKqHvG4","outputId":"c965a03c-5d9e-4051-9b33-330b063afe5f","execution":{"iopub.status.busy":"2023-06-15T09:14:49.821393Z","iopub.execute_input":"2023-06-15T09:14:49.821747Z","iopub.status.idle":"2023-06-15T09:14:50.023787Z","shell.execute_reply.started":"2023-06-15T09:14:49.821712Z","shell.execute_reply":"2023-06-15T09:14:50.022779Z"},"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        self.downsampling_path = nn.ModuleList([\n            Downsampling(in_channels, hid_channels, norm=False),\n            Downsampling(hid_channels, hid_channels*2),\n            Downsampling(hid_channels*2, hid_channels*4),\n            Downsampling(hid_channels*4, hid_channels*8),\n            Downsampling(hid_channels*8, hid_channels*8),\n            Downsampling(hid_channels*8, hid_channels*8),\n            Downsampling(hid_channels*8, hid_channels*8),\n            Downsampling(hid_channels*8, hid_channels*8, norm=False),\n        ])\n\n        self.upsampling_path = nn.ModuleList([\n            Upsampling(hid_channels*8, hid_channels*8, dropout=True),\n            Upsampling(hid_channels*16, hid_channels*8, dropout=True),\n            Upsampling(hid_channels*16, hid_channels*8, dropout=True),\n            Upsampling(hid_channels*16, hid_channels*8),\n            Upsampling(hid_channels*16, hid_channels*4),\n            Upsampling(hid_channels*8, hid_channels*2),\n            Upsampling(hid_channels*4, hid_channels),\n        ])\n\n        self.feature_block = nn.Sequential(\n            nn.ConvTranspose2d(hid_channels*2, out_channels, kernel_size=4, stride=2, padding=1),\n            nn.Tanh(),\n        )\n\n        self.feature_block.apply(init_weights)\n\n    def forward(self, x):\n        skips = []\n        for down in self.downsampling_path:\n            x = down(x)\n            skips.append(x)\n        skips = reversed(skips[:-1])\n\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","executionInfo":{"status":"ok","timestamp":1685394090858,"user_tz":420,"elapsed":565,"user":{"displayName":"Begona Garcia Malaxechebarria","userId":"06396578984430635223"}},"execution":{"iopub.status.busy":"2023-06-15T09:14:50.025427Z","iopub.execute_input":"2023-06-15T09:14:50.026057Z","iopub.status.idle":"2023-06-15T09:14:50.038284Z","shell.execute_reply.started":"2023-06-15T09:14:50.026003Z","shell.execute_reply":"2023-06-15T09:14:50.037189Z"},"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":2584,"status":"ok","timestamp":1685394103629,"user":{"displayName":"Begona Garcia Malaxechebarria","userId":"06396578984430635223"},"user_tz":420},"id":"7tZqDUCMKME8","outputId":"0676218e-b977-4257-9223-3f0913ff8751","execution":{"iopub.status.busy":"2023-06-15T09:14:50.041164Z","iopub.execute_input":"2023-06-15T09:14:50.041819Z","iopub.status.idle":"2023-06-15T09:14:52.01242Z","shell.execute_reply.started":"2023-06-15T09:14:50.041793Z","shell.execute_reply":"2023-06-15T09:14:52.011406Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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","execution":{"iopub.status.busy":"2023-06-15T09:15:21.533951Z","iopub.execute_input":"2023-06-15T09:15:21.534352Z","iopub.status.idle":"2023-06-15T09:15:21.542986Z","shell.execute_reply.started":"2023-06-15T09:15:21.534322Z","shell.execute_reply":"2023-06-15T09:15:21.542057Z"},"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","execution":{"iopub.status.busy":"2023-06-15T09:15:24.928984Z","iopub.execute_input":"2023-06-15T09:15:24.929359Z","iopub.status.idle":"2023-06-15T09:15:25.34632Z","shell.execute_reply.started":"2023-06-15T09:15:24.929331Z","shell.execute_reply":"2023-06-15T09:15:25.345275Z"},"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\nhid_channels = 64\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","execution":{"iopub.status.busy":"2023-06-15T09:14:52.549329Z","iopub.execute_input":"2023-06-15T09:14:52.549916Z","iopub.status.idle":"2023-06-15T09:14:52.843433Z","shell.execute_reply.started":"2023-06-15T09:14:52.549882Z","shell.execute_reply":"2023-06-15T09:14:52.842499Z"},"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":"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    ):\n        super().__init__()\n        self.save_hyperparameters(ignore=[\"optimizer\"])\n        self.optimizer = optimizer\n        self.automatic_optimization = False\n        self.step_num = 0\n        self.losses = {\n            \"gen_loss\": [],\n            \"disc_loss_M\": [],\n            \"disc_loss_P\": [],\n        }\n\n        self.monet_generator = monet_gen if monet_gen else UNetGenerator(hid_channels, in_channels, out_channels)\n        self.photo_generator = photo_gen if photo_gen else UNetGenerator(hid_channels, in_channels, out_channels)\n        self.monet_discriminator = monet_disc if monet_disc else Discriminator(hid_channels)\n        self.photo_discriminator = photo_disc if photo_disc else Discriminator(hid_channels)\n\n    def get_lr_scheduler(self, optimizer):\n        def lr_lambda(epoch):\n            len_decay_phase = self.hparams.num_epochs - self.hparams.decay_epochs + 1.0\n            curr_decay_step = max(0, epoch - self.hparams.decay_epochs + 1.0)\n            val = 1.0 - curr_decay_step / len_decay_phase\n            return max(0.0, val)\n\n        return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lr_lambda)\n\n    def forward(self, img):\n        return self.monet_generator(img)\n\n    def configure_optimizers(self):\n        opt_config = {\n            \"lr\": self.hparams.lr,\n            \"betas\": self.hparams.betas,\n        }\n        opt_gen = self.optimizer(\n            list(self.monet_generator.parameters()) + list(self.photo_generator.parameters()),\n            **opt_config,\n        )\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        schedulers = [self.get_lr_scheduler(opt) for opt in optimizers]\n        return optimizers, schedulers\n\n    def training_step(self, batch, batch_idx):\n        real_monet = batch[\"monet\"]\n        real_photo = batch[\"photo\"]\n        gen_optimizer, disc_optimizer = self.optimizers()\n\n        fake_monet = self.monet_generator(real_photo)\n        fake_photo = self.photo_generator(real_monet)\n\n        reconstructed_monet = self.monet_generator(fake_photo)\n        reconstructed_photo = self.photo_generator(fake_monet)\n\n        monet_identity = self.monet_generator(real_monet)\n        photo_identity = self.photo_generator(real_photo)\n\n        aug_real_monet = aug_fn(real_monet)\n        aug_fake_monet = aug_fn(fake_monet)\n\n        fake_monet_discrim_score = self.monet_discriminator(aug_fake_monet)\n        fake_photo_discrim_score = self.photo_discriminator(fake_photo)\n        real_monet_discrim_score = self.monet_discriminator(aug_real_monet)\n        real_photo_discrim_score = self.photo_discriminator(real_photo)\n\n        self.toggle_optimizer(gen_optimizer)\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,\n                                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(gen_optimizer)\n\n        self.toggle_optimizer(disc_optimizer)\n        monet_disc_loss = get_disc_loss(real_monet_discrim_score, self.monet_discriminator(aug_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(disc_optimizer)\n\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        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    def validation_step(self, batch, batch_idx):\n        self.display_results(batch, batch_idx, \"validate\")\n\n    def test_step(self, batch, batch_idx):\n        self.display_results(batch, batch_idx, \"test\")\n\n    def predict_step(self, batch, batch_idx):\n        return self(batch)\n\n    def display_results(self, batch, batch_idx, stage):\n        real_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        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        self.log(\"lr\", self.hparams.lr, on_step=False, on_epoch=True, prog_bar=True)\n        curr_lr = self.lr_schedulers()[0].get_last_lr()[0]\n        self.step = 0\n\n    def on_train_epoch_end(self):\n        for sch in self.lr_schedulers():\n            sch.step()\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        monet_gen_path = checkpoints + \"/monet_generator\" + str(self.current_epoch) + \".pth\"\n        torch.save(self.monet_generator.state_dict(), monet_gen_path)\n        photo_gen_path = checkpoints + \"/photo_generator\" + str(self.current_epoch) + \".pth\"\n        torch.save(self.photo_generator.state_dict(), photo_gen_path)\n        monet_disc_path = checkpoints + \"/monet_discriminator\" + str(self.current_epoch ) + \".pth\"\n        torch.save(self.monet_discriminator.state_dict(), monet_disc_path)\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):\n        predictions = self.trainer.predict_loop.predictions\n        print(f\"Done predicting\")\n","metadata":{"id":"VIZx1_P6hreW","execution":{"iopub.status.busy":"2023-06-15T09:15:34.944575Z","iopub.execute_input":"2023-06-15T09:15:34.94545Z","iopub.status.idle":"2023-06-15T09:15:34.973834Z","shell.execute_reply.started":"2023-06-15T09:15:34.945407Z","shell.execute_reply":"2023-06-15T09:15:34.972906Z"},"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":{"execution":{"iopub.status.busy":"2023-06-15T09:15:47.504325Z","iopub.execute_input":"2023-06-15T09:15:47.504699Z","iopub.status.idle":"2023-06-15T09:15:47.511441Z","shell.execute_reply.started":"2023-06-15T09:15:47.504671Z","shell.execute_reply":"2023-06-15T09:15:47.510476Z"},"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\": 21,\n}","metadata":{"id":"QT_DFMILg-aS","execution":{"iopub.status.busy":"2023-06-15T09:15:51.160175Z","iopub.execute_input":"2023-06-15T09:15:51.160564Z","iopub.status.idle":"2023-06-15T09:15:51.169102Z","shell.execute_reply.started":"2023-06-15T09:15:51.160534Z","shell.execute_reply":"2023-06-15T09:15:51.167881Z"},"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","execution":{"iopub.status.busy":"2023-06-15T09:35:17.806433Z","iopub.execute_input":"2023-06-15T09:35:17.806841Z","iopub.status.idle":"2023-06-15T09:35:17.814049Z","shell.execute_reply.started":"2023-06-15T09:35:17.806811Z","shell.execute_reply":"2023-06-15T09:35:17.813129Z"},"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","execution":{"iopub.status.busy":"2023-06-15T09:35:21.242374Z","iopub.execute_input":"2023-06-15T09:35:21.242747Z","iopub.status.idle":"2023-06-15T09:35:23.899045Z","shell.execute_reply.started":"2023-06-15T09:35:21.242717Z","shell.execute_reply":"2023-06-15T09:35:23.897271Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip uninstall deepspeed -y\nDS_BUILD_FUSED_ADAM=1 \n!pip install deepspeed ","metadata":{"execution":{"iopub.status.busy":"2023-06-15T09:33:13.755137Z","iopub.execute_input":"2023-06-15T09:33:13.755565Z","iopub.status.idle":"2023-06-15T09:33:27.577649Z","shell.execute_reply.started":"2023-06-15T09:33:13.755533Z","shell.execute_reply":"2023-06-15T09:33:27.576445Z"},"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_discriminator(example_monets[0])\nreal_monet_score2 = 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_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":57356,"status":"ok","timestamp":1684110025023,"user":{"displayName":"Garrett Devereux","userId":"06903598260790252274"},"user_tz":420},"id":"y0OqwP7xrNfA","outputId":"10444ba0-dfac-4c59-8286-47c427226c86","execution":{"iopub.status.busy":"2023-06-15T09:14:56.238273Z","iopub.status.idle":"2023-06-15T09:14:56.239467Z","shell.execute_reply.started":"2023-06-15T09:14:56.239187Z","shell.execute_reply":"2023-06-15T09:14:56.239213Z"},"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":{"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","execution":{"iopub.status.busy":"2023-06-15T09:14:56.240724Z","iopub.status.idle":"2023-06-15T09:14:56.24553Z","shell.execute_reply.started":"2023-06-15T09:14:56.24532Z","shell.execute_reply":"2023-06-15T09:14:56.245341Z"},"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-06-15T09:14:56.246857Z","iopub.status.idle":"2023-06-15T09:14:56.247677Z","shell.execute_reply.started":"2023-06-15T09:14:56.247383Z","shell.execute_reply":"2023-06-15T09:14:56.247409Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"checkpoints = \"/kaggle/working/\"\n\nmonet_gen_path = checkpoints + \"/monet_generator_diffaug\" + \".pth\"\ntorch.save(model.monet_generator.state_dict(), monet_gen_path)\n\nphoto_gen_path = checkpoints + \"/photo_generator_diffaug\" + \".pth\"\ntorch.save(model.photo_generator.state_dict(), photo_gen_path)\n\nmonet_disc_path = checkpoints + \"/monet_discriminator_diffaug\" + \".pth\"\ntorch.save(model.monet_discriminator.state_dict(), monet_disc_path)\n\nphoto_disc_path = checkpoints + \"/photo_discriminator_diffaug\" + \".pth\"\ntorch.save(model.photo_discriminator.state_dict(), photo_disc_path)","metadata":{"execution":{"iopub.status.busy":"2023-06-15T09:14:56.249588Z","iopub.status.idle":"2023-06-15T09:14:56.250455Z","shell.execute_reply.started":"2023-06-15T09:14:56.25018Z","shell.execute_reply":"2023-06-15T09:14:56.250205Z"},"trusted":true},"execution_count":null,"outputs":[]}]}