{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":21755,"databundleVersionId":1475600,"sourceType":"competition"}],"dockerImageVersionId":30698,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport os\nimport cv2\nimport random\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv\nimport matplotlib.pyplot as plt\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\n\nfrom torch.utils.data import Dataset, DataLoader\nfrom tqdm import tqdm\nfrom collections import defaultdict\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\n# for dirname, _, filenames in os.walk('/kaggle/input'):\n#     for filename in filenames:\n#         print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session\nplt.style.use(\"ggplot\")","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2024-05-25T11:27:30.345198Z","iopub.execute_input":"2024-05-25T11:27:30.345554Z","iopub.status.idle":"2024-05-25T11:27:33.502145Z","shell.execute_reply.started":"2024-05-25T11:27:30.345526Z","shell.execute_reply":"2024-05-25T11:27:33.501244Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def seed_everything(seed=1234):\n    random.seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.backends.cudnn.deterministic = True\nseed_everything()","metadata":{"execution":{"iopub.status.busy":"2024-05-25T11:27:34.434962Z","iopub.execute_input":"2024-05-25T11:27:34.435480Z","iopub.status.idle":"2024-05-25T11:27:34.443989Z","shell.execute_reply.started":"2024-05-25T11:27:34.435452Z","shell.execute_reply":"2024-05-25T11:27:34.443100Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## A PyTorch implementation of the CycleGAN based on the following: https://www.kaggle.com/code/ttymonkey/cyclegan-starter\n\n## GitHub Link: https://github.com/ADulian/CycleGAN/tree/main\nYou can find this implementation on my GitHub account where I'll be adding more features to it based on the original paper.\n\nIn addition, I developed a very simple tool that allows to easily explore the output of both Generators and Discriminators live and interactively. There is also a very simple notebook that you can easily run to play with that very tool. Here's a little demo\nhttps://github.com/ADulian/CycleGAN/blob/main/figures/cycle_gan_explorer.gif","metadata":{}},{"cell_type":"markdown","source":"# 1. Data Check","metadata":{}},{"cell_type":"code","source":"# Check the root\nroot_path = \"/kaggle/input/gan-getting-started\"\nos.listdir(root_path)","metadata":{"execution":{"iopub.status.busy":"2024-05-17T10:34:46.316715Z","iopub.execute_input":"2024-05-17T10:34:46.317196Z","iopub.status.idle":"2024-05-17T10:34:46.325811Z","shell.execute_reply.started":"2024-05-17T10:34:46.317167Z","shell.execute_reply":"2024-05-17T10:34:46.324867Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Reading img in rgb\nread_img = lambda path: cv2.cvtColor(cv2.imread(path), cv2.COLOR_BGR2RGB)","metadata":{"execution":{"iopub.status.busy":"2024-05-17T10:34:46.993926Z","iopub.execute_input":"2024-05-17T10:34:46.994275Z","iopub.status.idle":"2024-05-17T10:34:46.998651Z","shell.execute_reply.started":"2024-05-17T10:34:46.994247Z","shell.execute_reply":"2024-05-17T10:34:46.997721Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Load a sample image of monet and photo\ndata_path = f\"{root_path}/photo_jpg\"\nsample_photo = read_img(os.path.join(data_path, os.listdir(data_path)[0]))\n\ndata_path = f\"{root_path}/monet_jpg\"\nsample_monet = read_img(os.path.join(data_path, os.listdir(data_path)[0]))","metadata":{"execution":{"iopub.status.busy":"2024-05-17T10:34:47.158510Z","iopub.execute_input":"2024-05-17T10:34:47.158810Z","iopub.status.idle":"2024-05-17T10:34:47.518177Z","shell.execute_reply.started":"2024-05-17T10:34:47.158785Z","shell.execute_reply":"2024-05-17T10:34:47.517262Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"sample_photo.shape","metadata":{"execution":{"iopub.status.busy":"2024-05-17T10:34:47.519666Z","iopub.execute_input":"2024-05-17T10:34:47.519963Z","iopub.status.idle":"2024-05-17T10:34:47.525920Z","shell.execute_reply.started":"2024-05-17T10:34:47.519939Z","shell.execute_reply":"2024-05-17T10:34:47.524976Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"sample_photo.min(), sample_photo.max(), sample_photo.dtype","metadata":{"execution":{"iopub.status.busy":"2024-05-17T10:34:47.526913Z","iopub.execute_input":"2024-05-17T10:34:47.527154Z","iopub.status.idle":"2024-05-17T10:34:47.537723Z","shell.execute_reply.started":"2024-05-17T10:34:47.527133Z","shell.execute_reply":"2024-05-17T10:34:47.536593Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Photo\nplt.subplot(121)\nplt.title(\"Photo\")\nplt.imshow(sample_photo)  \n\n# Monet\nplt.subplot(122)\nplt.title(\"Photo\")\nplt.imshow(sample_monet)  ","metadata":{"execution":{"iopub.status.busy":"2024-05-17T10:34:47.615189Z","iopub.execute_input":"2024-05-17T10:34:47.615500Z","iopub.status.idle":"2024-05-17T10:34:48.161178Z","shell.execute_reply.started":"2024-05-17T10:34:47.615477Z","shell.execute_reply":"2024-05-17T10:34:48.160228Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# 2. Build The Model","metadata":{}},{"cell_type":"code","source":"# img params for testing\nIMG_H, IMG_W, IMG_C = 256, 256, 3","metadata":{"execution":{"iopub.status.busy":"2024-05-17T10:34:51.344716Z","iopub.execute_input":"2024-05-17T10:34:51.345128Z","iopub.status.idle":"2024-05-17T10:34:51.349375Z","shell.execute_reply.started":"2024-05-17T10:34:51.345103Z","shell.execute_reply":"2024-05-17T10:34:51.348476Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 2.1 Basic Blocks\n\nImplements 2 basic blocks\n- A Downsample block that aims to downsample input by 2 \n    - Input -> Conv2D -> (OPT) GroupNorm (InstanceNorm) -> LeakyReLU\n- An Upsample block that aims to upsample input by 2\n    - Input -> ConvTranspose2d -> GroupNorm (InstanceNorm) -> (OPT) Dropout -> ReLU","metadata":{}},{"cell_type":"code","source":"# --------------------------------------------------------------------------------\nclass Downsample(nn.Module):\n    \"\"\" A simple convolutional block that downsamples feature maps by 2 (stride=2)\n\n    For example\n        input_tensor [256, 256] -> output_tensor [128, 128]\n\n    Padding = same\n    Bias = False (although this will be checked)\n\n    Model\n        Conv2d -> (opt) InstanceNorm2d -> LeakyReLU\n\n    \"\"\"\n\n    # --------------------------------------------------------------------------------\n    def __init__(self,\n                 in_channels,\n                 out_channels,\n                 kernel_size,\n                 norm=True):\n\n        \"\"\" Init downsample module\n\n        ---\n        Parameters\n            in_channels: Number of input channels\n            out_channels: Number of output channels\n            kernel_size: Kernel size\n            norm: Whether to apply normalisation layer\n        \"\"\"\n\n        super().__init__()\n\n        # --- Layers\n        # Convolution Layer\n        self.conv = nn.Conv2d(in_channels=in_channels,\n                              out_channels=out_channels,\n                              kernel_size=kernel_size,\n                              stride=2,  # Downsample by 2\n                              padding=(kernel_size - 1) // 2,  # Same padding\n                              bias=False)  # ToDo: Check if this should be conditioned on norm\n\n        # Normalisation layer (equivalent of GroupNorm where num_groups=out_channels\n        if norm:\n            self.norm = nn.InstanceNorm2d(num_features=out_channels)\n        else:\n            self.norm = None\n\n        # Activation layer\n        self.act = nn.LeakyReLU()\n\n    # --------------------------------------------------------------------------------\n    def forward(self,\n                x):\n\n        \"\"\" Downsample input tensor\n\n        ---\n        Parameters\n            x: Input tensor\n\n        ---\n        Returns\n            out: x downsampled by 2\n\n        \"\"\"\n\n        out = self.conv(x)\n\n        if self.norm:\n            out = self.norm(out)\n\n        out = self.act(out)\n\n        return out\n","metadata":{"execution":{"iopub.status.busy":"2024-05-17T10:34:51.351095Z","iopub.execute_input":"2024-05-17T10:34:51.351366Z","iopub.status.idle":"2024-05-17T10:34:51.361472Z","shell.execute_reply.started":"2024-05-17T10:34:51.351343Z","shell.execute_reply":"2024-05-17T10:34:51.360648Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#  Check block\nx = torch.randn(1,IMG_C, IMG_H, IMG_W)\n\n# Downsample by 2 and Keep the same number of channels\nout = Downsample(3, 3, 3)(x)\n\nassert IMG_C == out.shape[1] \nassert IMG_H == out.shape[2] * 2\nassert IMG_W == out.shape[3] * 2","metadata":{"execution":{"iopub.status.busy":"2024-05-17T10:34:51.362489Z","iopub.execute_input":"2024-05-17T10:34:51.362782Z","iopub.status.idle":"2024-05-17T10:34:51.475927Z","shell.execute_reply.started":"2024-05-17T10:34:51.362759Z","shell.execute_reply":"2024-05-17T10:34:51.474904Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# --------------------------------------------------------------------------------\nclass Upsample(nn.Module):\n    \"\"\" A simple convolutional block that upsamples feature maps by 2 (kernel_size=4, stride=2)\n\n        For example\n            input_tensor [128, 128] -> output_tensor [256, 256]\n\n        Padding = same\n        Bias = False\n        Dropout = 0.5\n\n        Model\n            Conv2d -> InstanceNorm2d -> (opt) Dropout -> ReLU\n\n        \"\"\"\n\n    # --------------------------------------------------------------------------------\n    def __init__(self,\n                 in_channels,\n                 out_channels,\n                 kernel_size=4,\n                 dropout=True):\n        \"\"\" Init upsample module\n\n        ---\n        Parameters\n            in_channels: Number of input channels\n            out_channels: Number of output channels\n            kernel_size: Kernel size\n            dropout: Whether to apply dropout\n        \"\"\"\n\n        super().__init__()\n\n        # --- Layers\n        # Transpose Convolution Layer\n        self.conv = nn.ConvTranspose2d(in_channels=in_channels,\n                                       out_channels=out_channels,\n                                       kernel_size=kernel_size,\n                                       stride=2,  # Upsample by 2\n                                       padding=(kernel_size - 1) // 2,  # Same padding\n                                       bias=False)\n\n        # Normalisation layer (equivalent of GroupNorm where num_groups=out_channels\n        self.norm = nn.InstanceNorm2d(num_features=out_channels)\n\n        # Dropout\n        if dropout:\n            self.dropout = nn.Dropout(0.5)\n        else:\n            self.dropout = None\n\n        # Activation layer\n        self.act = nn.ReLU()\n\n    # --------------------------------------------------------------------------------\n    def forward(self,\n                x):\n        \"\"\" Upsample input tensor\n\n        ---\n        Parameters\n            x: Input tensor\n\n        ---\n        Returns\n            out: x upsampled by 2\n\n        \"\"\"\n\n        out = self.conv(x)\n        out = self.norm(out)\n\n        if self.dropout:\n            out = self.dropout(out)\n\n        out = self.act(out)\n\n        return out","metadata":{"execution":{"iopub.status.busy":"2024-05-17T10:34:51.477431Z","iopub.execute_input":"2024-05-17T10:34:51.478023Z","iopub.status.idle":"2024-05-17T10:34:51.492884Z","shell.execute_reply.started":"2024-05-17T10:34:51.477989Z","shell.execute_reply":"2024-05-17T10:34:51.490997Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#  Check block\nx = torch.randn(1,IMG_C, IMG_H, IMG_W)\n\n# Upsample by 2 and Keep the same number of channels\nout = Upsample(3, 3, 4)(x)\n\nassert IMG_C == out.shape[1] \nassert IMG_H == out.shape[2] // 2\nassert IMG_W == out.shape[3] // 2","metadata":{"execution":{"iopub.status.busy":"2024-05-17T10:34:51.494368Z","iopub.execute_input":"2024-05-17T10:34:51.495054Z","iopub.status.idle":"2024-05-17T10:34:51.592954Z","shell.execute_reply.started":"2024-05-17T10:34:51.495023Z","shell.execute_reply":"2024-05-17T10:34:51.591995Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 2.2 Generator","metadata":{}},{"cell_type":"code","source":"def conv_weight_init(m):\n    \"\"\" Initialise Conv2D and ConvTranspose2D with N(0, 0.02)\n    \"\"\"\n    if any(isinstance(m, _m) for _m in [nn.Conv2d, nn.ConvTranspose2d]):\n        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)\n","metadata":{"execution":{"iopub.status.busy":"2024-05-17T10:34:51.594235Z","iopub.execute_input":"2024-05-17T10:34:51.594710Z","iopub.status.idle":"2024-05-17T10:34:51.600519Z","shell.execute_reply.started":"2024-05-17T10:34:51.594655Z","shell.execute_reply":"2024-05-17T10:34:51.599548Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"A simple Generator network made out of an Encoder and Decoder with skip connections just like in U-Net (I believe)\n- Encoder\n    - 7 Downsample blocks that take input image and produce an output of shape [512, IMG_H//128, IMG_W//128]\n- Decoder\n    - 6 Upsample blocks that the latent features and produce an output of shape [64, LATENT_H * 64, LATENT_W * 64]\n- Final layer\n    - A final ConvTranspose2D layer that takes the output of Decoder and upsamples it once again so that the final output is of shape [3, IMG_H, IMG_W]","metadata":{}},{"cell_type":"code","source":"# --------------------------------------------------------------------------------\nclass Generator(nn.Module):\n    \"\"\" A Generator Network\n\n    ---\n    Structure\n        Img -> Encoder -> Decoder -> Img_Pred\n\n        Encoder\n            7 x Downsample blocks\n            Output -> Latent Tensor [512, IMG_H / 128, IMG_W / 128]\n        Decoder\n            6 x Upsample blocks\n            Output -> Upsampled Latent Tensor [64, IMG_H / 2, IMG_W / 2]\n        Final Upsample\n            ConvTranspose2D -> TanH\n            Output -> Target Image [IMG_C, IMG_H, IMG_W]\n\n    \"\"\"\n\n    # --------------------------------------------------------------------------------\n    def __init__(self):\n        \"\"\" Init CycleGAN's Generator\n\n        \"\"\"\n        super().__init__()\n\n        # --- Encoder and Decoder\n        self.encoder = self._init_encoder()\n        self.decoder = self._init_decoder()\n\n        # Final Upsample\n        self.out = nn.ConvTranspose2d(in_channels=128,\n                                      out_channels=3,\n                                      kernel_size=4,\n                                      stride=2, padding=1)\n\n        # Maybe softsign could be better than TanH\n        self.act = nn.Tanh()\n\n        # Initialise conv wights with N(0, 0.02)\n        self.apply(conv_weight_init)\n\n    # --------------------------------------------------------------------------------\n    def forward(self,\n                x):\n        \"\"\" Generate x_hat\n\n        ---\n        Parameters\n            x: Input tensor (img)\n\n        ---\n        Returns\n            torch.Tensor: Output tensor (img pred)\n\n        \"\"\"\n\n        # Encode to latent space / 128 #512\n        skips = []  # Skip connections U-Net Style\n        for layer in self.encoder:\n            x = layer(x)\n            skips.append(x)\n\n        # Skip last one, bottom bit, latent space\n        skips = reversed(skips[:-1])\n\n        # Decode form latent space *64 #128\n        for layer, skip in zip(self.decoder, skips):\n            x = layer(x)\n            x = torch.cat((x, skip), dim=1)\n\n        # Upsample so that out.shape == x.shape\n        out = self.act(self.out(x))\n\n        return out\n\n    # --------------------------------------------------------------------------------\n    def _init_encoder(self):\n        \"\"\" A Sequential Encoder with 7 Downsample Blocks each downsampling by 2\n\n        ---\n        Returns\n            nn.Sequential: Encoder\n\n        \"\"\"\n\n        # Settings\n        kernel_size = 4\n        in_channels = 64\n\n        # Final layer will have 512 feature maps and it's final output size is x / 128\n        out_channels = [128, 256, 512, 512, 512, 512]\n\n        # First layer doesn't use norm so just add it now\n        encoder = [Downsample(in_channels=3, out_channels=in_channels,\n                              kernel_size=kernel_size, norm=False)]\n\n        for out_ch in out_channels:\n            # Add layer\n            encoder.append(Downsample(in_channels=in_channels, out_channels=out_ch,\n                                      kernel_size=kernel_size))\n\n            # Update in_channels\n            in_channels = out_ch\n\n        return nn.Sequential(*encoder)\n\n    # --------------------------------------------------------------------------------\n    def _init_decoder(self):\n        \"\"\" A Sequential Decoder with 6 Upsample Blocks each upscaling by 2\n\n        ---\n        Returns\n            nn.Sequential: Decoder\n\n        \"\"\"\n\n        # Settings\n        kernel_size = 4\n        in_channels = [512, 1024, 1024, 1024, 512, 256]\n\n        # Final layer will have the same number of channels as the output of first downsample from encoder\n        out_channels = [512, 512, 512, 256, 128, 64]\n        dropout = [True, True, True, False, False, False]\n\n        decoder = [Upsample(in_channels=in_ch, out_channels=out_ch,\n                            kernel_size=kernel_size, dropout=drop)\n\n                   for in_ch, out_ch, drop in zip(in_channels, out_channels, dropout)]\n\n        return nn.Sequential(*decoder)","metadata":{"execution":{"iopub.status.busy":"2024-05-17T10:34:51.603763Z","iopub.execute_input":"2024-05-17T10:34:51.604224Z","iopub.status.idle":"2024-05-17T10:34:51.618750Z","shell.execute_reply.started":"2024-05-17T10:34:51.604189Z","shell.execute_reply":"2024-05-17T10:34:51.617877Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Test Generator\ng = Generator()\nx = torch.randn(1,3,256,256)\n\nenc_out = g.encoder(x)\ngen_out = g(x)\n\n# Downsample by 128\nassert enc_out.shape[1] == 512\nassert enc_out.shape[2] == x.shape[2] // 128\nassert enc_out.shape[3] == x.shape[3] // 128\n\n# Upsample to the same shape\nassert gen_out.shape[1] == x.shape[1]\nassert gen_out.shape[2] == x.shape[2]\nassert gen_out.shape[3] == x.shape[3]","metadata":{"execution":{"iopub.status.busy":"2024-05-17T10:34:51.619955Z","iopub.execute_input":"2024-05-17T10:34:51.620291Z","iopub.status.idle":"2024-05-17T10:34:52.596203Z","shell.execute_reply.started":"2024-05-17T10:34:51.620262Z","shell.execute_reply":"2024-05-17T10:34:52.595241Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Check weights if initialised correctly \nconvolutions = []\nfor _, m in g.named_modules():\n    if isinstance(m, nn.Conv2d):\n        convolutions.append(m.weight.flatten().clone().detach().numpy())\n\n# Here I expect values to be close to 0 and 0.02\nconvolutions = np.concatenate(convolutions)\nconvolutions.mean(), convolutions.std()","metadata":{"execution":{"iopub.status.busy":"2024-05-17T10:34:52.597651Z","iopub.execute_input":"2024-05-17T10:34:52.598207Z","iopub.status.idle":"2024-05-17T10:34:52.697221Z","shell.execute_reply.started":"2024-05-17T10:34:52.598167Z","shell.execute_reply":"2024-05-17T10:34:52.696220Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 2.3 Discriminator\nA simple Network that takes an input image of shape [3, 256, 256] and produces an output map [1, 30, 30] ","metadata":{}},{"cell_type":"code","source":"# --------------------------------------------------------------------------------\nclass DiscriminatorHistory:\n    \"\"\" Keep a history of fake samples that Discriminator can use to improve its robustness\n    \"\"\"\n\n    # --------------------------------------------------------------------------------\n    def __init__(self,\n                 buffer_size=50):\n        \"\"\" Init History module\n\n        ---\n        Parameters\n            buffer_size: Size of the buffer, as per paper the defualt is 50\n\n        The paper says \"We keep an image buffer that stores the 50 previously created images\",\n        however, what's worth noting is that the authors used a batch size of 1 for training,\n        meaning that the network did 50 updates before sampling from history. Worth rethinking\n        whether the buffer_size should be dependant on the batch_size\n\n        \"\"\"\n\n        self.buffer_size = buffer_size\n\n        # Empty Tensor\n        self.buffer = torch.Tensor()\n\n\n    # --------------------------------------------------------------------------------\n    def __call__(self,\n                 x):\n        \"\"\" Forward\n\n        Sample randomly from history buffer and add new sampels\n\n        ---\n        Parameters\n            x: Input Tensor\n\n        ---\n        Returns\n            Tensor with mixed samples from input and history\n\n        \"\"\"\n\n        # Make sure that the buffer is on the same device as x\n        if self.buffer.device != x.device:\n            self.buffer = self.buffer.to(x.device)\n\n        # Check if buffer is filled\n        if len(self.buffer) >= self.buffer_size:\n            # Sample from the buffer\n            sample_size = max(len(x) // 2, 1)\n            history_indices = list(range(0, self.buffer_size))\n            history_get_indices = random.sample(history_indices, k=sample_size)\n            x1 = self.buffer[history_get_indices]\n\n            if sample_size == 1:\n                return x1\n\n            # Sample from the input\n            x_indices = list(range(0, len(x)))\n            x_get_indices = random.sample(x_indices, k=sample_size)\n            x_set_indices = list(set(x_indices) - set(x_get_indices))\n            x2 = x[x_get_indices]\n\n            # Replace buffer with other set of samples\n            self.buffer[history_get_indices] = x[x_set_indices]\n\n            # Concat\n            x = torch.cat((x1, x2), 0)\n\n        # Fill the buffer\n        else:\n            self.buffer = torch.cat((self.buffer, x), 0)\n\n        # Return\n        return x\n","metadata":{"execution":{"iopub.status.busy":"2024-05-17T10:35:28.564590Z","iopub.execute_input":"2024-05-17T10:35:28.564960Z","iopub.status.idle":"2024-05-17T10:35:28.575482Z","shell.execute_reply.started":"2024-05-17T10:35:28.564932Z","shell.execute_reply":"2024-05-17T10:35:28.574605Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# --------------------------------------------------------------------------------\nclass Discriminator(nn.Module):\n    \"\"\" A Discriminator Network\n\n    ---\n    Structure\n        Img -> Discriminator -> Grid of predictions\n\n    \"\"\"\n\n    # --------------------------------------------------------------------------------\n    def __init__(self):\n        \"\"\" Init CycleGAN's Discriminator\n\n        \"\"\"\n\n        super().__init__()\n\n        # --- Default settings\n        base_channels = 64\n        kernel_size = 4\n\n        # --- History\n        self.history = DiscriminatorHistory()\n\n        # --- Discriminator network\n        self.discriminator = nn.Sequential(\n            Downsample(in_channels=3, out_channels=base_channels,\n                       kernel_size=kernel_size, norm=False),\n            Downsample(in_channels=base_channels, out_channels=base_channels * 2,\n                       kernel_size=kernel_size),\n            Downsample(in_channels=base_channels * 2, out_channels=base_channels * 4,\n                       kernel_size=kernel_size),\n            nn.ZeroPad2d(padding=(0, 2, 0, 2)),\n            nn.Conv2d(in_channels=base_channels * 4, out_channels=base_channels * 8,\n                      kernel_size=4, stride=1, bias=False),\n            nn.InstanceNorm2d(num_features=base_channels * 8),\n            nn.ZeroPad2d(padding=(0, 2, 0, 2)),\n            nn.Conv2d(in_channels=base_channels * 8, out_channels=1,\n                      kernel_size=4, stride=1),\n            #             nn.Sigmoid()\n        )\n\n        # Initialise conv wights with N(0, 0.02)\n        self.apply(conv_weight_init)\n\n    # --------------------------------------------------------------------------------\n    def forward(self,\n                x,\n                sample_history=False):\n        \"\"\" Discriminate between real/fake img\n\n        ---\n        Parameters\n            x: Input tensor (img)\n            sample_history: If true it will mix x with samples from buffer\n\n        ---\n        Returns\n            torch.Tensor: Output tensor (pred grid)\n\n        \"\"\"\n\n        # Update x based on history\n        if sample_history:\n            x = self.history(x)\n\n\n        return self.discriminator(x)","metadata":{"execution":{"iopub.status.busy":"2024-05-17T10:35:48.183094Z","iopub.execute_input":"2024-05-17T10:35:48.183471Z","iopub.status.idle":"2024-05-17T10:35:48.194073Z","shell.execute_reply.started":"2024-05-17T10:35:48.183442Z","shell.execute_reply":"2024-05-17T10:35:48.192945Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Test Discriminator\nd = Discriminator()\nx = torch.randn(1,3,256,256)\n\nassert d(x).shape[1] == 1","metadata":{"execution":{"iopub.status.busy":"2024-05-17T10:35:49.642090Z","iopub.execute_input":"2024-05-17T10:35:49.642457Z","iopub.status.idle":"2024-05-17T10:35:49.755296Z","shell.execute_reply.started":"2024-05-17T10:35:49.642428Z","shell.execute_reply":"2024-05-17T10:35:49.754219Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 2.4 Test with images","metadata":{}},{"cell_type":"code","source":"# To Torch Tensor\ntorch_sample = torch.tensor(sample_photo, dtype=torch.float32)\n\n# -1 to 1\ntorch_sample = (torch_sample / 255.0 - 0.5) / 0.5\n\n# H,W,C to C,H,W\ntorch_sample = torch_sample.permute(2,0,1)\n\n# Add Batch Dim 1,C,H,W\ntorch_sample = torch_sample.unsqueeze(dim=0)","metadata":{"execution":{"iopub.status.busy":"2024-05-17T10:35:52.808641Z","iopub.execute_input":"2024-05-17T10:35:52.809004Z","iopub.status.idle":"2024-05-17T10:35:52.815802Z","shell.execute_reply.started":"2024-05-17T10:35:52.808976Z","shell.execute_reply":"2024-05-17T10:35:52.814750Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"monet_generator = Generator()","metadata":{"execution":{"iopub.status.busy":"2024-05-17T10:35:57.505780Z","iopub.execute_input":"2024-05-17T10:35:57.506123Z","iopub.status.idle":"2024-05-17T10:35:58.255737Z","shell.execute_reply.started":"2024-05-17T10:35:57.506099Z","shell.execute_reply":"2024-05-17T10:35:58.254802Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"with torch.inference_mode():\n    monet_generator.eval()\n    to_monet = monet_generator(torch_sample).detach()\n    to_monet.shape","metadata":{"execution":{"iopub.status.busy":"2024-05-17T10:35:58.257519Z","iopub.execute_input":"2024-05-17T10:35:58.257907Z","iopub.status.idle":"2024-05-17T10:35:58.364656Z","shell.execute_reply.started":"2024-05-17T10:35:58.257875Z","shell.execute_reply":"2024-05-17T10:35:58.363703Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plt.subplot(1,2,1)\nplt.title(\"Original Photo\")\nplt.imshow(torch_sample[0].permute(1,2,0) * 0.5 + 0.5)\n\nplt.subplot(1,2,2)\nplt.title(\"Monet Photo\")\nplt.imshow(to_monet[0].permute(1,2,0) * 0.5 + 0.5)","metadata":{"execution":{"iopub.status.busy":"2024-05-17T10:35:58.365943Z","iopub.execute_input":"2024-05-17T10:35:58.366309Z","iopub.status.idle":"2024-05-17T10:35:59.071415Z","shell.execute_reply.started":"2024-05-17T10:35:58.366272Z","shell.execute_reply":"2024-05-17T10:35:59.070474Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 2.5 CycleGAN\n\nCycleGAN is built of 4 neural nets\n\n- 2 Generators\n    - Monet Generator - To go from Photo to Monet (G)\n    - Photo Generator - To got from Monet to Photo (F)\n- 2 Discriminators\n    - Monet Discriminator - To say whether input is a real/fake monet (Dy)\n    - Photo Discriminator - To say whether input is a real/fake photo (Dx)","metadata":{}},{"cell_type":"code","source":"# Define CycleGAN\n# --------------------------------------------------------------------------------\nclass CycleGAN(nn.Module):\n    \n    # --------------------------------------------------------------------------------\n    def __init__(self,\n                 lr=2e-4,\n                 lambda_cycle=10):  # Not sure what that is yet\n        \"\"\" Cycle GAN model\n        \n        ---\n        Structure:\n            - 2 Generators\n                - Monet Generator - To go from Photo to Monet (G)\n                - Photo Generator - To got from Monet to Photo (F)\n                \n            - 2 Discriminators\n                - Monet Discriminator - To say whether input is a real/fake monet (Dy)\n                - Photo Discriminator - To say whether input is a real/fake photo (Dx)\n\n                \n        \"\"\"\n        super().__init__()\n        \n        # Define Generators and Discriminators\n        self.gen_monet = Generator()\n        self.optim_gen_monet = optim.Adam(self.gen_monet.parameters(),\n                                          lr=lr,\n                                          betas=(0.5, 0.999))\n        self.gen_photo = Generator()\n        self.optim_gen_photo = optim.Adam(self.gen_photo.parameters(),\n                                          lr=lr,\n                                          betas=(0.5, 0.999))\n        \n        self.disc_monet = Discriminator()\n        self.optim_disc_monet = optim.Adam(self.disc_monet.parameters(),\n                                          lr=lr,\n                                          betas=(0.5, 0.999))\n        self.disc_photo = Discriminator()\n        self.optim_disc_photo = optim.Adam(self.disc_photo.parameters(),\n                                          lr=lr,\n                                          betas=(0.5, 0.999))\n        \n        # Loss functions\n        self.loss_l2 = nn.MSELoss()\n        self.loss_l1 = nn.L1Loss()\n        \n        # Attribs\n        self.lambda_cycle = lambda_cycle\n    \n    # --------------------------------------------------------------------------------\n    def forward(self, x):\n        \"\"\" Main point of the model is to go from photo -> monet so forward does just that\n        \n        Expects a single tensor (photo data)\n            x: [btach, 3, 256, 256] - Real photo\n        \"\"\"\n        \n        return self.gen_monet(x)\n    \n    # --------------------------------------------------------------------------------\n    def forward_apply_style(self,\n                            real_A,\n                            style=\"monet\"):\n        \"\"\" Apply one of the styles either monet or photo\n\n        In addition the function returns more than just styled photo e.g.\n        Disciminators on Fake, Real and Cycled Photo\n\n        ---\n        Parameters\n            real_A: A batch of real samples from domain A being transfered to domain B\n            style: Domain B, Applies either monet or photo style\n\n        ---\n        Returns\n            dict: Dictionary of outputs from model\n\n        \"\"\"\n\n        with torch.inference_mode(mode=True):\n            self.eval()\n\n            # Set networks and Domain\n            if style == \"monet\":\n                gen_A = self.gen_photo\n                gen_B = self.gen_monet\n\n                disc_A = self.disc_photo\n                disc_B = self.disc_monet\n            elif style == \"photo\":\n                gen_A = self.gen_photo\n                gen_B = self.gen_monet\n\n                disc_A = self.disc_photo\n                disc_B = self.disc_monet\n            else:\n                raise ValueError(f\"Style can be either monet or photo, given: {style}\")\n\n            # ---\n            # Forward, notation uses P=Photo and M=Monet to easier demonstrate the flow\n            # x_P -> G_M -> \\hat{x}_M\n            fake_B = gen_B(real_A)\n\n            # \\hat{x}_M -> G_P -> \\hat{x}_P\n            cycled_A = gen_A(fake_B)\n\n            # \\hat{x}_M -> D_M -> \\hat{y}_M\n            disc_fake_B = disc_B(fake_B)\n\n            # x_P -> D_P -> y_P\n            disc_real_A = disc_A(real_A)\n\n            # \\hat{x}_P -> D_M -> \\hat{y}_P\n            disc_cycled_A = disc_A(cycled_A)\n\n        return {\"real_A\" : real_A,\n                \"gen_fake_B\" : fake_B,\n                \"gen_cycled_A\" : cycled_A,\n                \"disc_fake_B\" : disc_fake_B,\n                \"disc_real_A\" : disc_real_A,\n                \"disc_cycled_A\" : disc_cycled_A}\n\n    # --------------------------------------------------------------------------------\n    def forward_step(self, x):\n        \"\"\" A single forward step for the model\n        \n        Forward -> Loss -> Backward -> Optimise\n        \n        Expects a tuple of 2 tensors\n            (tensor, tensor): [batch, 3, 256, 256] - Real monet and photo\n        \"\"\"\n        \n        # Unpack data\n        x_monet, x_photo = x\n        \n        # Zero Grad Optimisers\n        self.optim_zero_grad()\n        \n        # --- Generate fake data\n        x_fake_monet = self.gen_monet(x_photo)  # Real Photo -> Gm -> Fake Monet\n        x_fake_photo = self.gen_photo(x_monet)  # Real Monet -> Gp -> Fake Photo\n        \n        # --- Discriminators\n        # Monet\n        d_real_monet = self.disc_monet(x_monet)  # How real does this real monet seems\n        d_fake_monet = self.disc_monet(x_fake_monet.detach(), sample_history=True)  # How fake does this fake seems\n        \n        # Photo\n        d_real_photo = self.disc_photo(x_photo)\n        d_fake_photo = self.disc_photo(x_fake_photo.detach(), sample_history=True)\n        \n        # Discriminators Loss\n        d_loss_monet = self.disc_loss(x_real=d_real_monet,\n                                      x_fake=d_fake_monet)\n        \n        d_loss_photo = self.disc_loss(x_real=d_real_photo,\n                                      x_fake=d_fake_photo)\n        \n        d_loss = d_loss_monet + d_loss_photo\n        \n        # Backward\n        d_loss.backward()\n        \n        # Optimise\n        self.optim_disc_monet.step()\n        self.optim_disc_photo.step()\n        \n        # --- Generators\n        # Adversarial\n        x_adv_monet = self.disc_monet(x_fake_monet)  # How real does this fake monet seems\n        x_adv_photo = self.disc_photo(x_fake_photo)  \n        \n        # Cycle\n        x_cycled_monet = self.gen_monet(x_fake_photo)  # Fake Photo -> Gm -> Cycled Monet\n        x_cycled_photo = self.gen_photo(x_fake_monet)  # Fake Monet -> Gp -> Cycled Photo\n        \n        # Adversarial Loss\n        g_loss_adv_monet = self.loss_l2(torch.ones_like(x_adv_monet), x_adv_monet)\n        g_loss_adv_photo = self.loss_l2(torch.ones_like(x_adv_photo), x_adv_photo)\n        \n        # Cycle Loss\n        g_loss_cycle_monet = self.loss_l1(x_monet, x_cycled_monet) * self.lambda_cycle\n        g_loss_cycle_photo = self.loss_l1(x_photo, x_cycled_photo) * self.lambda_cycle\n        \n        # Identity Loss\n#         g_loss_idt_monet = self.loss_l1(x_monet, x_fake_monet) * self.lambda_cycle\n#         g_loss_idt_photo = self.loss_l1(x_photo, x_fake_photo) * self.lambda_cycle\n        \n        # Total Loss\n        g_loss_monet = g_loss_adv_monet + g_loss_cycle_monet #+ g_loss_idt_monet\n        g_loss_photo = g_loss_adv_photo + g_loss_cycle_photo #+ g_loss_idt_photo\n        g_loss = g_loss_monet + g_loss_photo\n        \n        # Backward\n        g_loss.backward()\n        \n        # Optimise\n        self.optim_gen_monet.step()\n        self.optim_gen_photo.step()\n        \n        # Get Losses\n        return {\n            \"g_monet\" : g_loss_monet,\n            \"g_photo\" : g_loss_photo,\n            \"d_monet\" : d_loss_monet,\n            \"d_photo\" : d_loss_photo\n        }\n    \n    # --------------------------------------------------------------------------------\n    def disc_loss(self,\n                  x_real,\n                  x_fake):\n        \"\"\" Compute loss for discriminator\n        \n        Loss = 1/2(Real loss + Fake loss)\n        \n        Expects 2 tensor\n            x_real: [batch, 3, 256, 256] - A real photo\n            x_fake: [batch, 3, 256, 256] - A fake photo from generator\n        \n        \"\"\"\n        \n        loss_real = self.loss_l2(torch.ones_like(x_real), x_real)\n        loss_fake = self.loss_l2(torch.zeros_like(x_fake), x_fake)\n        \n        return (loss_real + loss_fake) * 0.5\n    \n    # --------------------------------------------------------------------------------\n    def optim_zero_grad(self):\n        \"\"\" Apply zero grad on all optimisers\n        \"\"\"\n        \n        # Generators\n        self.optim_gen_monet.zero_grad()\n        self.optim_gen_photo.zero_grad()\n        \n        # Discriminators\n        self.optim_disc_monet.zero_grad()\n        self.optim_disc_photo.zero_grad()","metadata":{"execution":{"iopub.status.busy":"2024-05-17T10:36:20.534990Z","iopub.execute_input":"2024-05-17T10:36:20.535358Z","iopub.status.idle":"2024-05-17T10:36:20.561279Z","shell.execute_reply.started":"2024-05-17T10:36:20.535330Z","shell.execute_reply":"2024-05-17T10:36:20.560145Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"cycle_gan = CycleGAN()","metadata":{"execution":{"iopub.status.busy":"2024-05-17T10:36:21.488448Z","iopub.execute_input":"2024-05-17T10:36:21.488916Z","iopub.status.idle":"2024-05-17T10:36:23.119122Z","shell.execute_reply.started":"2024-05-17T10:36:21.488884Z","shell.execute_reply":"2024-05-17T10:36:23.118312Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"x_real = torch.ones(1, 3, 256, 256)\nx_fake = torch.zeros_like(x_real)\ncycle_gan.forward_step((x_real, x_fake))","metadata":{"execution":{"iopub.status.busy":"2024-05-17T10:36:23.120535Z","iopub.execute_input":"2024-05-17T10:36:23.120835Z","iopub.status.idle":"2024-05-17T10:36:26.138466Z","shell.execute_reply.started":"2024-05-17T10:36:23.120810Z","shell.execute_reply":"2024-05-17T10:36:26.137268Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# 3. Overfit on a Single Batch","metadata":{}},{"cell_type":"markdown","source":"## 3.1 Overfit","metadata":{}},{"cell_type":"code","source":"# Helpers\n\n# Np -> Tensor -> [batch, c, h, w] -> [-1, 1]\nto_tensor = lambda x: torch.tensor(x, dtype=torch.float32).permute(2,0,1).unsqueeze(0) / 127.5 - 1.0\n\n# Tensor -> [h, w, c] -> Np -> [0, 255]\nfrom_tensor = lambda x: ((x.squeeze().permute(1,2,0).numpy() + 1.0) * 127.5).astype(np.uint8)","metadata":{"execution":{"iopub.status.busy":"2024-05-17T10:36:26.147020Z","iopub.execute_input":"2024-05-17T10:36:26.147647Z","iopub.status.idle":"2024-05-17T10:36:26.156395Z","shell.execute_reply.started":"2024-05-17T10:36:26.147612Z","shell.execute_reply":"2024-05-17T10:36:26.155435Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Set device\ndevice = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")","metadata":{"execution":{"iopub.status.busy":"2024-05-17T10:36:27.581521Z","iopub.execute_input":"2024-05-17T10:36:27.582327Z","iopub.status.idle":"2024-05-17T10:36:27.586661Z","shell.execute_reply.started":"2024-05-17T10:36:27.582293Z","shell.execute_reply":"2024-05-17T10:36:27.585741Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Pair of images to torch tensor and [-1,1]\ntensor_monet = to_tensor(sample_monet).to(device)\ntensor_photo = to_tensor(sample_photo).to(device)\n\ntensor_monet.min(), tensor_monet.max(), tensor_monet.shape, tensor_monet.device","metadata":{"execution":{"iopub.status.busy":"2024-05-17T10:36:27.749583Z","iopub.execute_input":"2024-05-17T10:36:27.750150Z","iopub.status.idle":"2024-05-17T10:36:28.102489Z","shell.execute_reply.started":"2024-05-17T10:36:27.750124Z","shell.execute_reply":"2024-05-17T10:36:28.101600Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Re-init model - SKIP THIS CELL TO CONTINUE TRAINING \ncycle_gan = CycleGAN(lr=1e-6).to(device)  # Keep LR very small for a single sample\n\n# Keep losses per epoch\nloss_epoch = defaultdict(list)","metadata":{"execution":{"iopub.status.busy":"2024-05-17T10:36:29.811707Z","iopub.execute_input":"2024-05-17T10:36:29.812336Z","iopub.status.idle":"2024-05-17T10:36:31.545834Z","shell.execute_reply.started":"2024-05-17T10:36:29.812296Z","shell.execute_reply":"2024-05-17T10:36:31.545027Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Training params\nNUM_EPOCHS = 1  # ToDo: Increase to test\n\n# Validate every  10% of epochs\nVALIDATE_EVERY = max(int(0.1 * NUM_EPOCHS), 1)\nPHOTO_TO_MONET = []\n\n# Train for N epochs\nwith tqdm(total=NUM_EPOCHS,\n          desc=\"Training\",\n          unit=\"Epoch\",\n          ascii=True, \n          colour=\"green\") as pbar:\n    \n    for i in range(1, NUM_EPOCHS + 1):\n\n        # Batch of data\n        loss = cycle_gan.forward_step((tensor_monet, tensor_photo))\n        \n        # Update dict\n        for k, v in loss.items():\n            loss_epoch[k].append(v.item())\n\n        # Bar Update\n        pbar.set_postfix({k:v.item() for k, v in loss.items()})\n        pbar.update()\n        \n        # Validate\n        if i % VALIDATE_EVERY == 0:\n            with torch.inference_mode(mode=True):\n                cycle_gan.eval()\n                monet_pred = cycle_gan(tensor_photo)\n            \n            PHOTO_TO_MONET.append(monet_pred)\n            cycle_gan.train()","metadata":{"execution":{"iopub.status.busy":"2024-05-17T10:36:47.582337Z","iopub.execute_input":"2024-05-17T10:36:47.582996Z","iopub.status.idle":"2024-05-17T10:36:48.629864Z","shell.execute_reply.started":"2024-05-17T10:36:47.582968Z","shell.execute_reply":"2024-05-17T10:36:48.628815Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 3.2 Visualise","metadata":{}},{"cell_type":"code","source":"# Quick visualisation of loss\nplt.figure(figsize=(10,8))\n\n# Generators\nplt.subplot(221)\nplt.plot(range(len(loss_epoch[\"g_monet\"])), loss_epoch[\"g_monet\"])\nplt.title(\"G Monet\")\n\nplt.subplot(222)\nplt.plot(range(len(loss_epoch[\"g_photo\"])), loss_epoch[\"g_photo\"])\nplt.title(\"G Photo\")\n\n# Discriminators\nplt.subplot(223)\nplt.plot(range(len(loss_epoch[\"d_monet\"])), loss_epoch[\"d_monet\"])\nplt.title(\"D Monet\")\n\nplt.subplot(224)\nplt.plot(range(len(loss_epoch[\"d_photo\"])), loss_epoch[\"d_photo\"])\nplt.title(\"D Photo\")\n\n\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2024-05-17T10:36:51.417210Z","iopub.execute_input":"2024-05-17T10:36:51.417558Z","iopub.status.idle":"2024-05-17T10:36:52.175656Z","shell.execute_reply.started":"2024-05-17T10:36:51.417534Z","shell.execute_reply":"2024-05-17T10:36:52.174650Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Visualise that sample\nwith torch.inference_mode(mode=True):\n    cycle_gan.eval()\n    monet_pred = cycle_gan(tensor_photo)\n\nplt.figure(figsize=(15,10))\n    \nplt.subplot(131)\nplt.title(\"Photo\")\nplt.imshow(from_tensor(tensor_photo.cpu()))\n\nplt.subplot(132)\nplt.title(\"Monet\")\nplt.imshow(sample_monet)\n\nplt.subplot(133)\nplt.title(\"Photo to Monet\")\nplt.imshow(from_tensor(monet_pred.cpu()))\n\n","metadata":{"execution":{"iopub.status.busy":"2024-05-17T10:36:55.852503Z","iopub.execute_input":"2024-05-17T10:36:55.853122Z","iopub.status.idle":"2024-05-17T10:36:56.804994Z","shell.execute_reply.started":"2024-05-17T10:36:55.853090Z","shell.execute_reply":"2024-05-17T10:36:56.804059Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Visualise All samples from Validation steps\n# Create a figure and axes\nfig, axes = plt.subplots(2, 5, figsize=(20, 10))\n\n# Flatten the axes array\naxes = axes.flatten()\n\n# Loop through images and plot them\nfor i, img in enumerate(PHOTO_TO_MONET):\n    axes[i].imshow(from_tensor(img.cpu()))\n    axes[i].axis('off')\n    axes[i].set_title(f\"Photo to Monet at: {(i+1) * 10}%\")\n\nplt.tight_layout()\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2024-05-17T10:36:57.369820Z","iopub.execute_input":"2024-05-17T10:36:57.370868Z","iopub.status.idle":"2024-05-17T10:36:58.988021Z","shell.execute_reply.started":"2024-05-17T10:36:57.370835Z","shell.execute_reply":"2024-05-17T10:36:58.987097Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# 4. Dataset ","metadata":{}},{"cell_type":"markdown","source":"A very simple PyTorch Dataset object\n\n`__getitem__` returns a photo between `[0, len(photos) -1]` based on given `idx` as well as a random monet photo as its pair\n\nOne epoch will loop over all photos whilst seeing monets several times (7k vs 300 images of monet)","metadata":{}},{"cell_type":"code","source":"class MonetDataset(Dataset):\n    def __init__(self, \n                 root_path):\n        \n        # Paths\n        self.root_path = root_path\n        self.photo_paths = os.listdir(f\"{root_path}/photo_jpg\")\n        self.monet_paths = os.listdir(f\"{root_path}/monet_jpg\")\n        \n    def __len__(self):\n        return len(self.photo_paths)\n    \n    def __getitem__(self, idx):\n        \"\"\" Get one of the photos and a random monet\n        \"\"\"\n        \n        # Get a photo path\n        path_photo = f\"{root_path}/photo_jpg/{self.photo_paths[idx]}\"\n        \n        # Get a random monet path\n        path_monet = f\"{root_path}/monet_jpg/{random.choices(self.monet_paths, k=1)[0]}\"\n        \n        # Load imgs and transform\n        x_photo = self.to_tensor(self.read_img(path_photo))\n        x_monet = self.to_tensor(self.read_img(path_monet))\n        \n        return x_monet, x_photo\n        \n    def to_tensor(self, x):\n        \"\"\" Transform uint8 image [0,255] to torch float32 [-1,1]\n        \"\"\"\n        return torch.tensor(x, dtype=torch.float32).permute(2,0,1) / 127.5 - 1.0\n        \n    def read_img(self, path):\n        \"\"\" Read img with cv2 and transform to RGB\n        \"\"\"\n        return cv2.cvtColor(cv2.imread(path), cv2.COLOR_BGR2RGB)        ","metadata":{"execution":{"iopub.status.busy":"2024-05-17T10:37:00.534237Z","iopub.execute_input":"2024-05-17T10:37:00.535200Z","iopub.status.idle":"2024-05-17T10:37:00.545066Z","shell.execute_reply.started":"2024-05-17T10:37:00.535165Z","shell.execute_reply":"2024-05-17T10:37:00.543778Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# 5. Training","metadata":{}},{"cell_type":"markdown","source":"## 5.1 Train the Model","metadata":{}},{"cell_type":"code","source":"# Training params\nNUM_EPOCHS = 15\nBATCH_SIZE = 16\n\n# Device \ndevice = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")\n\n# Keep losses per epoch\nloss_epoch = defaultdict(list)\n\n# Dataset and data loader\ndataset = MonetDataset(root_path=root_path)\ndata_loader = DataLoader(dataset=dataset, batch_size=BATCH_SIZE, shuffle=True, pin_memory=True, num_workers=4)\n\n# Model\ncycle_gan = CycleGAN(lr=2e-4,).to(device)","metadata":{"execution":{"iopub.status.busy":"2024-05-17T10:37:32.866070Z","iopub.execute_input":"2024-05-17T10:37:32.866833Z","iopub.status.idle":"2024-05-17T10:37:34.491751Z","shell.execute_reply.started":"2024-05-17T10:37:32.866803Z","shell.execute_reply":"2024-05-17T10:37:34.490972Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Train for N epochs\nfor i in range(1, NUM_EPOCHS + 1):  # Starting at one for tqdm\n    \n    # Tqdm\n    with tqdm(total=len(data_loader),\n          desc=f\"Epoch {i}/{NUM_EPOCHS}\",\n          ascii=True, \n          colour=\"green\") as pbar:\n        \n        # Per batch\n        loss_batch = defaultdict(list)\n        for x in data_loader:\n            x_monet = x[0].to(device)\n            x_photo = x[1].to(device)\n\n            # Batch of data\n            loss = cycle_gan.forward_step((x_monet, x_photo))\n\n            # Update dict\n            for k, v in loss.items():\n                loss_batch[k].append(v.item())\n\n            # Bar Update\n            pbar.set_postfix({k:v.item() for k, v in loss.items()})\n            pbar.update()\n            \n        # Update mean loss for epoch\n        for k, v in loss_batch.items():\n            loss_epoch[k].append(np.mean(v))","metadata":{"execution":{"iopub.status.busy":"2024-04-25T20:16:51.712208Z","iopub.execute_input":"2024-04-25T20:16:51.712512Z","iopub.status.idle":"2024-04-25T21:10:09.720493Z","shell.execute_reply.started":"2024-04-25T20:16:51.712487Z","shell.execute_reply":"2024-04-25T21:10:09.719257Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# 6. Submissions","metadata":{}},{"cell_type":"code","source":"import PIL\n! mkdir ../images","metadata":{"execution":{"iopub.status.busy":"2024-04-25T21:19:29.918958Z","iopub.execute_input":"2024-04-25T21:19:29.919744Z","iopub.status.idle":"2024-04-25T21:19:30.936655Z","shell.execute_reply.started":"2024-04-25T21:19:29.919708Z","shell.execute_reply":"2024-04-25T21:19:30.935324Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"i = 1\ndata_loader = DataLoader(dataset=dataset, batch_size=1)\nwith torch.inference_mode(mode=True):\n    cycle_gan.eval()\n    \n    for x in data_loader:\n        photo, _ = x\n        \n        pred = cycle_gan(photo.to(device)).cpu().squeeze().numpy()\n        pred = np.transpose(pred, (1,2,0))\n        \n        pred = (pred * 127.5 + 127.5).astype(np.uint8)\n        \n        pred_img = PIL.Image.fromarray(pred)\n        pred_img.save(\"../images/\" + str(i) + \".jpg\")\n        \n        i += 1","metadata":{"execution":{"iopub.status.busy":"2024-04-25T21:26:01.068581Z","iopub.execute_input":"2024-04-25T21:26:01.068946Z","iopub.status.idle":"2024-04-25T21:27:35.061142Z","shell.execute_reply.started":"2024-04-25T21:26:01.068916Z","shell.execute_reply":"2024-04-25T21:27:35.060110Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import shutil\nshutil.make_archive(\"/kaggle/working/images\", 'zip', \"/kaggle/images\")","metadata":{"execution":{"iopub.status.busy":"2024-04-25T21:27:35.238057Z","iopub.execute_input":"2024-04-25T21:27:35.238382Z","iopub.status.idle":"2024-04-25T21:27:38.537919Z","shell.execute_reply.started":"2024-04-25T21:27:35.238355Z","shell.execute_reply":"2024-04-25T21:27:38.536903Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# 7 Visualise results","metadata":{}},{"cell_type":"code","source":"# Quick visualisation of loss\nplt.figure(figsize=(10,8))\n\n# Generators\nplt.subplot(221)\nplt.plot(range(len(loss_epoch[\"g_monet\"])), loss_epoch[\"g_monet\"])\nplt.title(\"G Monet\")\n\nplt.subplot(222)\nplt.plot(range(len(loss_epoch[\"g_photo\"])), loss_epoch[\"g_photo\"])\nplt.title(\"G Photo\")\n\n# Discriminators\nplt.subplot(223)\nplt.plot(range(len(loss_epoch[\"d_monet\"])), loss_epoch[\"d_monet\"])\nplt.title(\"D Monet\")\n\nplt.subplot(224)\nplt.plot(range(len(loss_epoch[\"d_photo\"])), loss_epoch[\"d_photo\"])\nplt.title(\"D Photo\")\n\n\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2024-04-25T21:10:09.723019Z","iopub.execute_input":"2024-04-25T21:10:09.723392Z","iopub.status.idle":"2024-04-25T21:10:10.553424Z","shell.execute_reply.started":"2024-04-25T21:10:09.723356Z","shell.execute_reply":"2024-04-25T21:10:10.552496Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plt.figure(figsize=(15,10))\n\nplt.subplot(121)\nplt.title(\"Photo\")\nplt.imshow(from_tensor(tensor_photo.cpu()))\nplt.axis(\"off\")\n\nplt.subplot(122)\nplt.title(\"Monet\")\nplt.imshow(sample_monet)\nplt.axis(\"off\")","metadata":{"execution":{"iopub.status.busy":"2024-04-25T21:17:04.254359Z","iopub.execute_input":"2024-04-25T21:17:04.255048Z","iopub.status.idle":"2024-04-25T21:17:05.022943Z","shell.execute_reply.started":"2024-04-25T21:17:04.255016Z","shell.execute_reply":"2024-04-25T21:17:05.021874Z"},"trusted":true},"execution_count":null,"outputs":[]}]}