{"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":"## Import Necessary Libraries","metadata":{"id":"YYSTgmPRp8xm"}},{"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-15T12:50:10.013167Z","iopub.execute_input":"2023-06-15T12:50:10.013561Z","iopub.status.idle":"2023-06-15T12:50:32.552625Z","shell.execute_reply.started":"2023-06-15T12:50:10.013529Z","shell.execute_reply":"2023-06-15T12:50:32.551396Z"},"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":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-15T12:50:32.555264Z","iopub.execute_input":"2023-06-15T12:50:32.556022Z","iopub.status.idle":"2023-06-15T12:50:32.567009Z","shell.execute_reply.started":"2023-06-15T12:50:32.555961Z","shell.execute_reply":"2023-06-15T12:50:32.565944Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Load and Preprocess the Data","metadata":{"id":"5_LMm626Iz3r"}},{"cell_type":"markdown","source":"### Storing The Dataset","metadata":{"id":"NZOk1fMqKGz3"}},{"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-15T12:50:32.568550Z","iopub.execute_input":"2023-06-15T12:50:32.568996Z","iopub.status.idle":"2023-06-15T12:50:32.576487Z","shell.execute_reply.started":"2023-06-15T12:50:32.568952Z","shell.execute_reply":"2023-06-15T12:50:32.575456Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### DataLoaders and DataModules","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        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        )","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-15T12:50:32.580161Z","iopub.execute_input":"2023-06-15T12:50:32.580744Z","iopub.status.idle":"2023-06-15T12:50:32.591710Z","shell.execute_reply.started":"2023-06-15T12:50:32.580712Z","shell.execute_reply":"2023-06-15T12:50:32.590689Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Configuration Details","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","executionInfo":{"status":"ok","timestamp":1685393976138,"user_tz":420,"elapsed":19,"user":{"displayName":"Begona Garcia Malaxechebarria","userId":"06396578984430635223"}},"execution":{"iopub.status.busy":"2023-06-15T12:50:32.593381Z","iopub.execute_input":"2023-06-15T12:50:32.593713Z","iopub.status.idle":"2023-06-15T12:50:32.603312Z","shell.execute_reply.started":"2023-06-15T12:50:32.593683Z","shell.execute_reply":"2023-06-15T12:50:32.602378Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Understanding the Dataset","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":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-15T12:50:32.604926Z","iopub.execute_input":"2023-06-15T12:50:32.605534Z","iopub.status.idle":"2023-06-15T12:50:35.356267Z","shell.execute_reply.started":"2023-06-15T12:50:32.605501Z","shell.execute_reply":"2023-06-15T12:50:35.355040Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Augmentation","metadata":{}},{"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\")","metadata":{"execution":{"iopub.status.busy":"2023-06-15T12:50:35.358375Z","iopub.execute_input":"2023-06-15T12:50:35.358794Z","iopub.status.idle":"2023-06-15T12:50:35.371636Z","shell.execute_reply.started":"2023-06-15T12:50:35.358751Z","shell.execute_reply":"2023-06-15T12:50:35.369747Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Visualizing the Dataset","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-15T12:50:35.373351Z","iopub.execute_input":"2023-06-15T12:50:35.374030Z","iopub.status.idle":"2023-06-15T12:50:35.384678Z","shell.execute_reply.started":"2023-06-15T12:50:35.373995Z","shell.execute_reply":"2023-06-15T12:50:35.383724Z"},"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-15T12:50:35.386343Z","iopub.execute_input":"2023-06-15T12:50:35.386682Z","iopub.status.idle":"2023-06-15T12:50:36.835721Z","shell.execute_reply.started":"2023-06-15T12:50:35.386652Z","shell.execute_reply":"2023-06-15T12:50:36.834901Z"},"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-15T12:50:36.840112Z","iopub.execute_input":"2023-06-15T12:50:36.840684Z","iopub.status.idle":"2023-06-15T12:50:38.813203Z","shell.execute_reply.started":"2023-06-15T12:50:36.840649Z","shell.execute_reply":"2023-06-15T12:50:38.811904Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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-15T12:50:38.814709Z","iopub.execute_input":"2023-06-15T12:50:38.815751Z","iopub.status.idle":"2023-06-15T12:50:38.821867Z","shell.execute_reply.started":"2023-06-15T12:50:38.815716Z","shell.execute_reply":"2023-06-15T12:50:38.820883Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#### Downsampling Module","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","executionInfo":{"status":"ok","timestamp":1685394033859,"user_tz":420,"elapsed":59,"user":{"displayName":"Begona Garcia Malaxechebarria","userId":"06396578984430635223"}},"execution":{"iopub.status.busy":"2023-06-15T12:50:38.823749Z","iopub.execute_input":"2023-06-15T12:50:38.824359Z","iopub.status.idle":"2023-06-15T12:50:38.834383Z","shell.execute_reply.started":"2023-06-15T12:50:38.824303Z","shell.execute_reply":"2023-06-15T12:50:38.833398Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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-15T12:50:38.835920Z","iopub.execute_input":"2023-06-15T12:50:38.836513Z","iopub.status.idle":"2023-06-15T12:50:38.986522Z","shell.execute_reply.started":"2023-06-15T12:50:38.836480Z","shell.execute_reply":"2023-06-15T12:50:38.984710Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#### Upsampling Module","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-15T12:50:38.988181Z","iopub.execute_input":"2023-06-15T12:50:38.988542Z","iopub.status.idle":"2023-06-15T12:50:38.996462Z","shell.execute_reply.started":"2023-06-15T12:50:38.988509Z","shell.execute_reply":"2023-06-15T12:50:38.995265Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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-15T12:50:38.998545Z","iopub.execute_input":"2023-06-15T12:50:38.998889Z","iopub.status.idle":"2023-06-15T12:50:39.131236Z","shell.execute_reply.started":"2023-06-15T12:50:38.998856Z","shell.execute_reply":"2023-06-15T12:50:39.130189Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Build the Generator","metadata":{"id":"sUsii0Ihcz-J"}},{"cell_type":"code","source":"OUTPUT_CHANNELS = 3\n\ndef downsample(filters, size, apply_instancenorm=True):\n    initializer = tf.random_normal_initializer(0., 0.02)\n    gamma_init = tf.keras.initializers.RandomNormal(mean=0., stddev=0.02)\n\n    model_down = tf.keras.Sequential()\n    model_down.add(layers.Conv2D(filters, size, strides=2, padding='same',\n                                kernel_initializer=initializer, use_bias=False))\n\n    if apply_instancenorm:\n        model_down.add(tfa.layers.InstanceNormalization(gamma_initializer=gamma_init))\n\n    model_down.add(layers.LeakyReLU())\n\n    return model_down","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-15T12:50:39.132644Z","iopub.execute_input":"2023-06-15T12:50:39.133203Z","iopub.status.idle":"2023-06-15T12:50:39.140228Z","shell.execute_reply.started":"2023-06-15T12:50:39.133167Z","shell.execute_reply":"2023-06-15T12:50:39.139150Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def upsample(filters, size, apply_dropout=False):\n    initializer = tf.random_normal_initializer(0., 0.02)\n    gamma_init = tf.keras.initializers.RandomNormal(mean=0., stddev=0.02)\n\n    model_up = tf.keras.Sequential()\n    model_up.add(layers.Conv2DTranspose(filters, size, strides=2, padding='same',\n                                kernel_initializer=initializer, use_bias=False))\n\n    model_up.add(tfa.layers.InstanceNormalization(gamma_initializer=gamma_init))\n\n    if apply_dropout:\n        model_up.add(layers.Dropout(0.5))\n\n    model_up.add(layers.ReLU())\n\n    return model_up","metadata":{"execution":{"iopub.status.busy":"2023-06-15T12:50:39.141606Z","iopub.execute_input":"2023-06-15T12:50:39.142192Z","iopub.status.idle":"2023-06-15T12:50:39.154263Z","shell.execute_reply.started":"2023-06-15T12:50:39.142158Z","shell.execute_reply":"2023-06-15T12:50:39.153282Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def Generator():\n    inputs = layers.Input([256,256,3])\n\n    down_stack = [\n        downsample(64, 4, apply_instancenorm=False),\n        downsample(128, 4),\n        downsample(256, 4),\n        downsample(512, 4),\n        downsample(512, 4),\n        downsample(512, 4),\n        downsample(512, 4),\n        downsample(512, 4),\n    ]\n\n    up_stack = [\n        upsample(512, 4, apply_dropout=True),\n        upsample(512, 4, apply_dropout=True),\n        upsample(512, 4, apply_dropout=True),\n        upsample(512, 4),\n        upsample(256, 4),\n        upsample(128, 4),\n        upsample(64, 4),\n    ]\n\n    initializer = tf.random_normal_initializer(0., 0.02)\n    last = layers.Conv2DTranspose(OUTPUT_CHANNELS, 4, strides=2, padding='same',\n                                kernel_initializer=initializer, activation='tanh')\n\n    x = inputs\n\n    skips = []\n    for down in down_stack:\n        x = down(x)\n        skips.append(x)\n\n    skips = reversed(skips[:-1])\n    \n    for up, skip in zip(up_stack, skips):\n        x = up(x)\n        x = layers.Concatenate()([x, skip])\n\n    x = last(x)\n\n    return tf.keras.Model(inputs=inputs, outputs=x)","metadata":{"execution":{"iopub.status.busy":"2023-06-15T12:50:39.155826Z","iopub.execute_input":"2023-06-15T12:50:39.156397Z","iopub.status.idle":"2023-06-15T12:50:39.166700Z","shell.execute_reply.started":"2023-06-15T12:50:39.156366Z","shell.execute_reply":"2023-06-15T12:50:39.166097Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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-15T12:50:39.167892Z","iopub.execute_input":"2023-06-15T12:50:39.168471Z","iopub.status.idle":"2023-06-15T12:50:41.068105Z","shell.execute_reply.started":"2023-06-15T12:50:39.168440Z","shell.execute_reply":"2023-06-15T12:50:41.067105Z"},"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-15T12:50:41.069688Z","iopub.execute_input":"2023-06-15T12:50:41.070730Z","iopub.status.idle":"2023-06-15T12:50:41.079205Z","shell.execute_reply.started":"2023-06-15T12:50:41.070693Z","shell.execute_reply":"2023-06-15T12:50:41.078200Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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-15T12:50:41.080369Z","iopub.execute_input":"2023-06-15T12:50:41.080692Z","iopub.status.idle":"2023-06-15T12:50:41.513639Z","shell.execute_reply.started":"2023-06-15T12:50:41.080667Z","shell.execute_reply":"2023-06-15T12:50:41.511908Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Visualizing the Generators and Discriminators - Before Training","metadata":{"id":"fEW50f9ceNSP"}},{"cell_type":"code","source":"# Base Dimensions to use\nhid_channels = 64\nin_channels = 3\nout_channels = 3\n\n# Generator to transform Photos to Monet Paintings\nmonet_generator = UNetGenerator(hid_channels, in_channels, out_channels)\n# Generator to transform Monet Paintings to Photos\nphoto_generator = UNetGenerator(hid_channels, in_channels, out_channels)\n\n# Give a batch of example photos to each generator\nto_monet = monet_generator(example_photos)\nto_photos = photo_generator(example_monets)\n\nshow_img(example_photos, nrow=5, title=\"Original Photos\")\nshow_img(to_monet, nrow=5, title=\"Generated Monets\")\nshow_img(example_monets, nrow=5, title=\"Original Monets\")\nshow_img(to_photos, nrow=5, title=\"Generated Photos\")","metadata":{"executionInfo":{"elapsed":22353,"status":"ok","timestamp":1684109653977,"user":{"displayName":"Garrett Devereux","userId":"06903598260790252274"},"user_tz":420},"id":"3LvRnlEVeaPW","outputId":"eded2845-2aac-4294-a4e1-4c132ce8a5fd","execution":{"iopub.status.busy":"2023-06-15T12:50:41.516700Z","iopub.execute_input":"2023-06-15T12:50:41.517008Z","iopub.status.idle":"2023-06-15T12:50:47.102840Z","shell.execute_reply.started":"2023-06-15T12:50:41.516980Z","shell.execute_reply":"2023-06-15T12:50:47.101908Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Discriminator to identify real Monet Paintings\nmonet_discriminator = Discriminator(hid_channels)\n# Discriminator to identify real photos\nphoto_discriminator = Discriminator(hid_channels)\n\n# Give a real photo to each discriminator (first one photo shown above)\nreal_monet_scores = monet_discriminator(example_monets[0])\nreal_photo_scores = photo_discriminator(example_photos[0])\n\n# Sizes are randomly initialized\nprint(\"Monet Discriminator Size and Scores\")\nprint(real_monet_scores.shape)\nprint(real_monet_scores)\nprint()\nprint(\"Photo Discriminator Size and Scores\")\nprint(real_photo_scores.shape)\nprint(real_photo_scores)","metadata":{"executionInfo":{"elapsed":437,"status":"ok","timestamp":1684109665575,"user":{"displayName":"Garrett Devereux","userId":"06903598260790252274"},"user_tz":420},"id":"EqXpm9eaf0UY","outputId":"161f1c23-a8d7-466d-973b-69f9bbfa6f06","execution":{"iopub.status.busy":"2023-06-15T12:50:47.104180Z","iopub.execute_input":"2023-06-15T12:50:47.105080Z","iopub.status.idle":"2023-06-15T12:50:47.411619Z","shell.execute_reply.started":"2023-06-15T12:50:47.105047Z","shell.execute_reply":"2023-06-15T12:50:47.410667Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Build THE CycleGAN\n","metadata":{"id":"8HhJGCR4gXrl"}},{"cell_type":"markdown","source":"### Generator Loss","metadata":{"id":"PPjVeQXAhlim"}},{"cell_type":"code","source":"# How good are we at tricking the discriminator\ndef get_adv_loss(fake_image_discriminator_score):\n    # A perfect fake will have the discriminator produce all 1s\n    real_labels = torch.ones_like(fake_image_discriminator_score)\n    # Retrun mse between actual scores and 1s\n    return F.mse_loss(fake_image_discriminator_score, real_labels)","metadata":{"id":"YLh8EtJOktSQ","execution":{"iopub.status.busy":"2023-06-15T12:50:47.413229Z","iopub.execute_input":"2023-06-15T12:50:47.413892Z","iopub.status.idle":"2023-06-15T12:50:47.420407Z","shell.execute_reply.started":"2023-06-15T12:50:47.413857Z","shell.execute_reply":"2023-06-15T12:50:47.419379Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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-15T12:50:47.421928Z","iopub.execute_input":"2023-06-15T12:50:47.422396Z","iopub.status.idle":"2023-06-15T12:50:47.430170Z","shell.execute_reply.started":"2023-06-15T12:50:47.422343Z","shell.execute_reply":"2023-06-15T12:50:47.429298Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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 * F.l1_loss(identity_image, real)","metadata":{"id":"PBxx6ehUmlYH","execution":{"iopub.status.busy":"2023-06-15T12:50:47.431467Z","iopub.execute_input":"2023-06-15T12:50:47.431844Z","iopub.status.idle":"2023-06-15T12:50:47.439240Z","shell.execute_reply.started":"2023-06-15T12:50:47.431811Z","shell.execute_reply":"2023-06-15T12:50:47.438318Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Computes the Loss of BOTH Generators, given\n# - Original Photos from dataset, real_monet and real_photo\n# - Generated Fakes, fake_monet and fake_photo\n# - Reconstructed Photos, reconstructed_monet, reconstructed_photo\n# - Identities, monet_identity and photo_identity\n# - Discrim Scores of fake images, fake_monet_discrim_score, fake_photo_discrim_score\n# - Model Hyperparams lambda_cycle and lambda_idt\ndef get_gen_loss(real_monet, real_photo, fake_monet, fake_photo, reconstructed_monet, reconstructed_photo,\n                 monet_identity, photo_identity, fake_monet_discrim_score, fake_photo_discrim_score, lambda_cycle, lambda_idt):\n  \n\n    # STEP 1: Compute Adversarial Loss using fakes and discriminator scores\n    # Conceptually - how good the generatros are at tricking the discriminator\n    # Adv Loss of monet generator\n    adv_loss_photo_to_monet = get_adv_loss(fake_monet_discrim_score)\n    # Adv Loss of photo generator\n    adv_loss_monet_to_photo = get_adv_loss(fake_photo_discrim_score)\n    total_adv_loss = adv_loss_photo_to_monet + adv_loss_monet_to_photo\n            \n    # STEP 2: Compute Cycle Loss using\n    # Cycle loss of the original monet compared to reconstructed monet\n    cycle_loss_MPM = get_cycle_loss(real_monet, reconstructed_monet, lambda_cycle[0])\n    # Cycle loss of the original photo compared to reconstructed photo\n    cycle_loss_PMP = get_cycle_loss(real_photo, reconstructed_photo, lambda_cycle[1])\n    total_cycle_loss = cycle_loss_MPM + cycle_loss_PMP\n\n    # STEP 3: Compute Identity Loss using reals and identities\n    # Identity Loss of the real monet and monet identity\n    idt_loss_MM = get_idt_loss(real_monet, monet_identity, lambda_idt)\n    # Identity Loss of the real photo and identity photo\n    idt_loss_PP = get_idt_loss(real_photo, photo_identity, lambda_idt)\n    total_idt_loss = idt_loss_MM + idt_loss_PP\n            \n    # FINAL STEP: Sum losses to get total generator loss\n    gen_loss = total_adv_loss + total_idt_loss + total_cycle_loss\n    return gen_loss","metadata":{"id":"lkIdtRoDueDy","execution":{"iopub.status.busy":"2023-06-15T12:50:47.440891Z","iopub.execute_input":"2023-06-15T12:50:47.441306Z","iopub.status.idle":"2023-06-15T12:50:47.451809Z","shell.execute_reply.started":"2023-06-15T12:50:47.441274Z","shell.execute_reply":"2023-06-15T12:50:47.450942Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Discriminator Loss","metadata":{"id":"z0HIvxSzoCz9"}},{"cell_type":"code","source":"# Computes loss of Single Discriminator given fake and real discriminator scores\ndef get_disc_loss(real_discrim_score, fake_discrim_score):\n    # STEP 1: Calculate real discrim loss\n    # Our discriminator should output all 1s as this is a real image from the dataset\n    real_labels = torch.ones_like(real_discrim_score)\n    # Return mse between actual scores and 1s\n    real_loss = F.mse_loss(real_discrim_score, real_labels)\n        \n    # STEP 2: Calculate fake discrim loss\n    # Our discriminator should output all 0s as this is a generated image from the dataset\n    fake_labels = torch.zeros_like(fake_discrim_score)\n    # Return mse between actual scores and 1s\n    fake_loss = F.mse_loss(fake_discrim_score, fake_labels)\n        \n    # Return average of losses\n    disc_loss = (fake_loss + real_loss) * 0.5\n    return disc_loss","metadata":{"id":"CZMXvaBmn9Sg","execution":{"iopub.status.busy":"2023-06-15T12:50:47.458488Z","iopub.execute_input":"2023-06-15T12:50:47.458857Z","iopub.status.idle":"2023-06-15T12:50:47.465068Z","shell.execute_reply.started":"2023-06-15T12:50:47.458796Z","shell.execute_reply":"2023-06-15T12:50:47.464045Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### PyTorch Lightning Module","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    ):\n        super().__init__()\n        self.save_hyperparameters(ignore=[\"optimizer\"])\n        self.optimizer = optimizer\n        # Necessary to set as we are using multiple optimizers\n        self.automatic_optimization = False\n\n        # Step for Graphs\n        self.step_num = 0\n        \n        # INITIALIZE OUR GENERATORS AND DISCRIMINATORS!!!!\n        # Weights will automatically be initialized here\n\n        # Generators:\n        # Generator to transform Photos to Monet Paintings\n        if monet_gen == None:\n          self.monet_generator = UNetGenerator(hid_channels, in_channels, out_channels)\n        else:\n          self.monet_generator = monet_gen\n  \n        # Generator to transform Monet Paintings into Photos\n        if photo_gen == None:\n          self.photo_generator = UNetGenerator(hid_channels, in_channels, out_channels)\n        else:\n          self.photo_generator = photo_gen\n\n        # Discriminators\n        # Discriminator to identify real Monet Paintings\n        if monet_disc == None:\n          self.monet_discriminator = Discriminator(hid_channels)\n        else:\n          self.monet_discriminator = monet_disc\n\n        # Discriminator to identify real Photos\n        if photo_disc == None:\n          self.photo_discriminator = Discriminator(hid_channels)\n        else:\n          self.photo_discriminator = photo_disc\n\n        # Metrics for Loss graphs\n        self.losses = {\n            \"gen_loss\": [],\n            \"disc_loss_M\": [],\n            \"disc_loss_P\": [],\n        }\n        \n    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        # Set up optimizer for discriminators\n        opt_disc = self.optimizer(\n            list(self.monet_discriminator.parameters()) + list(self.photo_discriminator.parameters()),\n            **opt_config,\n        )\n        optimizers = [opt_gen, opt_disc]\n        # No schedulers for now\n        schedulers = [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        # Store the optimizers for the generators and discriminators\n        gen_optimizer, disc_optimizer = self.optimizers()\n\n        # STEP 1: Use Generators to produce fake images in both domains\n        # Fake monet generated from the real photo\n        fake_monet = self.monet_generator(real_photo)\n        # Fake photo generated from the real monet\n        fake_photo = self.photo_generator(real_monet)\n         \n        # STEP 2: Use generators and fakes to reconstruct original inputs\n        # Reconstruction of original monet painting\n        reconstructed_monet = self.monet_generator(fake_photo)\n        # Reconstruction of original photo\n        reconstructed_photo = self.photo_generator(fake_monet)\n\n        # STEP 3: Construct identities by feeding real inputs into the generators\n        # Monet identity: feed real monet into monet generator\n        monet_identity = self.monet_generator(real_monet)\n        # Photo identity: feed real photo into photo generator\n        photo_identity = self.photo_generator(real_photo)\n        \n        ## 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        # STEP 4: Compute Discriminator Scores to see how well our\n        # Discriminators differentiate real and fake images\n        # Scores of how good the fake generated monets are\n        fake_monet_discrim_score = self.monet_discriminator(aug_fake_monet)\n        # Scores of how good the fake generated photos are\n        fake_photo_discrim_score = self.photo_discriminator(fake_photo)\n        # Scores of how good the real monets are\n        real_monet_discrim_score = self.monet_discriminator(aug_real_monet)\n        # Scores of how good the real photos are\n        real_photo_discrim_score = self.photo_discriminator(real_photo)\n    \n        # STEP 5: Use Generator Loss to train Generators\n        # Standard pipeline: Loss -> zero_grad -> backward -> step\n        self.toggle_optimizer(gen_optimizer)\n        gen_loss = get_gen_loss(real_monet, real_photo, fake_monet, fake_photo, reconstructed_monet, reconstructed_photo,\n                 monet_identity, photo_identity, fake_monet_discrim_score, fake_photo_discrim_score, self.hparams.lambda_cycle, self.hparams.lambda_idt)      \n        gen_optimizer.zero_grad()\n        self.manual_backward(gen_loss)\n        gen_optimizer.step()\n        self.untoggle_optimizer(gen_optimizer)\n        \n        # STEP 5: Use Discriminator Loss to train Discriminators\n        # Standard pipeline: Losses -> zero_grad -> backward -> step\n        self.toggle_optimizer(disc_optimizer)\n        # Need to detatch otherwise backward is happening twice\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        # Done with training step\n        # Log metrics and losses\n        metrics = {\n            \"gen_loss\": gen_loss,\n            \"disc_loss_M\": monet_disc_loss,\n            \"disc_loss_P\": photo_disc_loss,\n        }\n        self.log_dict(metrics, on_step=False, on_epoch=True, prog_bar=True)\n\n\n        self.losses[\"gen_loss\"].append(float(gen_loss))\n        self.losses[\"disc_loss_M\"].append(float(monet_disc_loss))\n        self.losses[\"disc_loss_P\"].append(float(photo_disc_loss))\n\n        self.step += 1\n        \n        \n    # Output current progress of generator on a validation set\n    def validation_step(self, batch, batch_idx):\n        self.display_results(batch, batch_idx, \"validate\")\n    \n    # Output current progress of generator on a test set\n    def test_step(self, batch, batch_idx):\n        self.display_results(batch, batch_idx, \"test\")\n        \n    # Predict the results of a batch\n    def predict_step(self, batch, batch_idx):\n        return self(batch)\n    \n    # Function to Visualize Results\n    def display_results(self, batch, batch_idx, stage):\n        # Store batch of real input photos and corresponding fake generated monets\n        real_photos = batch\n        fake_monets = self(real_photos)\n        \n        if stage == \"validate\":\n            title = f\"Epoch {self.current_epoch+1}: Photo-to-Monet Translation\"\n        else:\n            title = f\"Sample {batch_idx+1}: Photo-to-Monet Translation\"\n\n        # Display current image results\n        show_img(\n            torch.cat([real_photos, fake_monets], dim=0),\n            nrow=len(real_photos),\n            title=title,\n        )\n    \n    def on_train_epoch_start(self):\n        # Update Progress Bar\n        self.log(\"lr\", self.hparams.lr, on_step=False, on_epoch=True, prog_bar=True)\n        curr_lr = self.lr_schedulers()[0].get_last_lr()[0]\n        # Reset Step\n        self.step = 0\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        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":{"id":"VIZx1_P6hreW","execution":{"iopub.status.busy":"2023-06-15T12:50:47.466576Z","iopub.execute_input":"2023-06-15T12:50:47.467556Z","iopub.status.idle":"2023-06-15T12:50:47.503018Z","shell.execute_reply.started":"2023-06-15T12:50:47.467525Z","shell.execute_reply":"2023-06-15T12:50:47.502192Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Training Configuration","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-15T12:50:47.504113Z","iopub.execute_input":"2023-06-15T12:50:47.505235Z","iopub.status.idle":"2023-06-15T12:50:47.516699Z","shell.execute_reply.started":"2023-06-15T12:50:47.505201Z","shell.execute_reply":"2023-06-15T12:50:47.515805Z"},"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-15T12:50:47.518185Z","iopub.execute_input":"2023-06-15T12:50:47.518609Z","iopub.status.idle":"2023-06-15T12:50:47.528487Z","shell.execute_reply.started":"2023-06-15T12:50:47.518577Z","shell.execute_reply":"2023-06-15T12:50:47.527590Z"},"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-15T12:50:47.530025Z","iopub.execute_input":"2023-06-15T12:50:47.530462Z","iopub.status.idle":"2023-06-15T12:50:47.539739Z","shell.execute_reply.started":"2023-06-15T12:50:47.530431Z","shell.execute_reply":"2023-06-15T12:50:47.538807Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## TRAIN THE CYLCEGAN","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-15T12:50:47.542875Z","iopub.execute_input":"2023-06-15T12:50:47.543202Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Post Training Visualization\n\nLet's visualize how the Loss of both of the Generators as well as the individual Discriminator Losses changed throughout the training process.","metadata":{"id":"Zi6IjSqZqSb_"}},{"cell_type":"code","source":"# Loss Plots\nfig, axes = plt.subplots(ncols=1, nrows=2, figsize=(18, 12), facecolor='w')\ntraining_steps = np.arange(len(model.losses[\"gen_loss\"])) / 4\n\n# Plot the Generator Loss\naxes[0].set_title(\"Generator Loss\")\naxes[0].plot(training_steps, model.losses[\"gen_loss\"], label='Generator Loss')\naxes[0].legend()\naxes[0].set_xlabel('Epochs')\n\n# Plot the Discriminator Losses\naxes[1].set_title(\"Discriminator Losses\")\naxes[1].plot(training_steps, model.losses[\"disc_loss_M\"], label='Monet Discriminator Loss')\naxes[1].plot(training_steps, model.losses[\"disc_loss_P\"], label='Photo Discriminator Loss')\naxes[1].legend()\naxes[1].set_xlabel('Epochs')\n\nplt.show()","metadata":{"executionInfo":{"elapsed":1713,"status":"ok","timestamp":1683872676459,"user":{"displayName":"Garrett Devereux","userId":"06903598260790252274"},"user_tz":420},"id":"VJE4FwVGF36B","outputId":"adb2d61d-0b69-4157-eea3-9eba2dc58810","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Visualizing Trained Generators and Discriminators","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","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","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","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":{"trusted":true},"execution_count":null,"outputs":[]}]}