{"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":30732,"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 numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\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\nimport os\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","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2024-07-16T20:41:14.274357Z","iopub.execute_input":"2024-07-16T20:41:14.275007Z","iopub.status.idle":"2024-07-16T20:41:14.280379Z","shell.execute_reply.started":"2024-07-16T20:41:14.274966Z","shell.execute_reply":"2024-07-16T20:41:14.279381Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Import Libraries and Modules","metadata":{}},{"cell_type":"code","source":"import torch #design and train neural networks\nfrom torch.utils.data import Dataset, DataLoader #customize dataset, feed data\nimport torchvision #image to tensors, resize and normalize images\nimport matplotlib.pyplot as plt #visualize data and nn training progress\nfrom PIL import Image # open load images","metadata":{"execution":{"iopub.status.busy":"2024-07-16T20:41:16.975196Z","iopub.execute_input":"2024-07-16T20:41:16.975835Z","iopub.status.idle":"2024-07-16T20:41:16.980397Z","shell.execute_reply.started":"2024-07-16T20:41:16.975803Z","shell.execute_reply":"2024-07-16T20:41:16.979470Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Create PyTorch Dataset Class","metadata":{}},{"cell_type":"code","source":"class TorchDataset(Dataset):\n    def __init__(self, data_path, transform=None):\n        self.data_path = data_path\n        self.transform = transform\n        \n    # enable len() func for Dataset class\n    def __len__(self):\n        return len(self.data_path)\n    \n    # fetch image by idx, open with PIL, apply transformation\n    def __getitem__(self, idx):\n        image_path = [os.path.join(self.data_path, _) for _ in os.listdir(self.data_path)]\n        index_path = image_path[idx]\n        image = Image.open(index_path)\n        image = self.transform(image)\n        return image ","metadata":{"execution":{"iopub.status.busy":"2024-07-16T20:41:19.265556Z","iopub.execute_input":"2024-07-16T20:41:19.266442Z","iopub.status.idle":"2024-07-16T20:41:19.272952Z","shell.execute_reply.started":"2024-07-16T20:41:19.266407Z","shell.execute_reply":"2024-07-16T20:41:19.272060Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Define Transformations \n* Resize image to 256*256 to unify input dimensions\n* Convert to tensors for neural networks\n* Normalize the tensor value to be between -1~1 by setting mean=0.5 sd=0.5 (across 3 RGB channels)\n\n","metadata":{}},{"cell_type":"code","source":"transforms = torchvision.transforms.Compose([\n    torchvision.transforms.Resize((256,256)),\n    torchvision.transforms.ToTensor(),\n    torchvision.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))\n])\n\n","metadata":{"execution":{"iopub.status.busy":"2024-07-16T20:41:22.570576Z","iopub.execute_input":"2024-07-16T20:41:22.570941Z","iopub.status.idle":"2024-07-16T20:41:22.576263Z","shell.execute_reply.started":"2024-07-16T20:41:22.570913Z","shell.execute_reply":"2024-07-16T20:41:22.575343Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# DataLoader and Exploratory data analysis","metadata":{}},{"cell_type":"code","source":"root_path = \"/kaggle/input/gan-getting-started\"\nprint(os.listdir(root_path))\ndata_path_monet = f\"{root_path}/monet_jpg\"\ndata_path_photo = f\"{root_path}/photo_jpg\"","metadata":{"execution":{"iopub.status.busy":"2024-07-16T20:41:26.081898Z","iopub.execute_input":"2024-07-16T20:41:26.082303Z","iopub.status.idle":"2024-07-16T20:41:26.088786Z","shell.execute_reply.started":"2024-07-16T20:41:26.082270Z","shell.execute_reply":"2024-07-16T20:41:26.087819Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"MonetDataset = TorchDataset(data_path=data_path_monet, transform=transforms)\nPhotoDataset = TorchDataset(data_path=data_path_photo, transform=transforms)\n\nMonetLoader = DataLoader(MonetDataset, batch_size=8, shuffle=True, num_workers=2)\nPhotoLoader = DataLoader(PhotoDataset, batch_size=8, shuffle=True, num_workers=2)","metadata":{"execution":{"iopub.status.busy":"2024-07-16T20:41:54.506621Z","iopub.execute_input":"2024-07-16T20:41:54.506998Z","iopub.status.idle":"2024-07-16T20:41:54.512600Z","shell.execute_reply.started":"2024-07-16T20:41:54.506969Z","shell.execute_reply":"2024-07-16T20:41:54.511689Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# fetch the next batch(32 images in tensor format) from the two DataLoaders \nmonets = next(iter(MonetLoader))\nphotos = next(iter(PhotoLoader))\nprint(monets.shape, photos.shape)\n","metadata":{"execution":{"iopub.status.busy":"2024-07-16T20:41:55.285890Z","iopub.execute_input":"2024-07-16T20:41:55.286274Z","iopub.status.idle":"2024-07-16T20:41:55.899683Z","shell.execute_reply.started":"2024-07-16T20:41:55.286244Z","shell.execute_reply":"2024-07-16T20:41:55.898514Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(monets[0].min(), monets[0].max())\nprint(photos[0].min(), photos[0].max())","metadata":{"execution":{"iopub.status.busy":"2024-07-16T20:41:55.980832Z","iopub.execute_input":"2024-07-16T20:41:55.981197Z","iopub.status.idle":"2024-07-16T20:41:55.989641Z","shell.execute_reply.started":"2024-07-16T20:41:55.981138Z","shell.execute_reply":"2024-07-16T20:41:55.988777Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# tansform data to be readable by plt.imshow\nmonets_imshow = []\nfor monet in monets:\n    monet = monet.permute(1, 2, 0) #tensor[C,H,W] -> imshow[H,W,C] height,weight,channel\n    monet = monet.numpy() # tensor -> numpy_arr readable by imshow\n    monet = monet * 0.5 + 0.5 # un-normalize the data\n    monets_imshow.append(monet)\n\nphotos_imshow = []    \nfor photo in photos:\n    photo = photo.permute(1, 2, 0) #tensor[C,H,W] -> imshow[H,W,C] height,weight,channel\n    photo = photo.numpy() # tensor -> numpy_arr readable by imshow\n    photo = photo * 0.5 + 0.5 # un-normalize the data\n    photos_imshow.append(photo)","metadata":{"execution":{"iopub.status.busy":"2024-07-16T20:41:57.862869Z","iopub.execute_input":"2024-07-16T20:41:57.863295Z","iopub.status.idle":"2024-07-16T20:41:57.882299Z","shell.execute_reply.started":"2024-07-16T20:41:57.863262Z","shell.execute_reply":"2024-07-16T20:41:57.881292Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(monets_imshow[3].shape)  # Check the shape of one of the images\nprint(monets_imshow[3].dtype) ","metadata":{"execution":{"iopub.status.busy":"2024-07-16T20:41:58.401226Z","iopub.execute_input":"2024-07-16T20:41:58.402135Z","iopub.status.idle":"2024-07-16T20:41:58.407470Z","shell.execute_reply.started":"2024-07-16T20:41:58.402079Z","shell.execute_reply":"2024-07-16T20:41:58.406556Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"\n# Monets\nfig1 = plt.figure(figsize=(15, 5)) \n\nplt.subplot(1, 4, 1)  \nplt.imshow(monets_imshow[0])\nplt.axis('off')  \n\nplt.subplot(1, 4, 2)  \nplt.imshow(monets_imshow[1])\nplt.axis('off')  \n\nplt.subplot(1, 4, 3)  \nplt.imshow(monets_imshow[2])\nplt.axis('off')  \n\nplt.subplot(1, 4, 4)  \nplt.imshow(monets_imshow[3])\nplt.axis('off')  \n\n\n# Add a title for the entire figure\nfig1.suptitle('Monets', fontsize=16)\n\n# Display all the plots\nplt.show()\n\n\n\n\n# Photos\nfig2 = plt.figure(figsize=(15, 5)) \n\nplt.subplot(1, 4, 1)  \nplt.imshow(photos_imshow[0])\nplt.axis('off')  \n\nplt.subplot(1, 4, 2)  \nplt.imshow(photos_imshow[1])\nplt.axis('off')  \n\nplt.subplot(1, 4, 3)  \nplt.imshow(photos_imshow[2])\nplt.axis('off')  \n\nplt.subplot(1, 4, 4)  \nplt.imshow(photos_imshow[3])\nplt.axis('off')  \n\n\n# Add a title for the entire figure\nfig2.suptitle('Photos', fontsize=16)\n\n# Display all the plots\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2024-07-16T20:41:59.606081Z","iopub.execute_input":"2024-07-16T20:41:59.606480Z","iopub.status.idle":"2024-07-16T20:42:00.375071Z","shell.execute_reply.started":"2024-07-16T20:41:59.606444Z","shell.execute_reply":"2024-07-16T20:42:00.374194Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# CycleGAN Model","metadata":{}},{"cell_type":"markdown","source":"### Define Block Class\nIn neural network architecture design, a \"block\" typically refers to a repeatable and modular unit of layers that perform a set of operations, which can be stacked or repeated with similar or varying parameters to build more complex architectures. The advantage of designing neural networks in blocks includes:\n* Reusability: Once defined, blocks can be reused within different parts of the network or in different network architectures.\n* Readability and Maintainability: Organizing code into logical units (blocks) makes it easier to manage and understand.\n* Modularity: Blocks can be designed, tested, and optimized independently.","metadata":{}},{"cell_type":"code","source":"import torch.nn as nn\n\nclass Block(nn.Module):\n    # define states (attributes) -> NN layers\n    def __init__(self, in_channels, out_channels, stride):\n        # inherent from the parent class -> nn.Module\n        super().__init__()\n        self.cnn = nn.Sequential(\n            # Convolution layer\n            nn.Conv2d(\n                in_channels,\n                out_channels,\n                kernel_size=4,\n                stride=stride,\n                padding=1,\n                bias=True,\n                padding_mode=\"reflect\", # reflect padding reduces artifact\n            ), \n            # Normalization layer -> Instance rather than Batch\n            # helps preserve artistic style integrity, \n            # regardless of the variability in content across different images in a batch\n            nn.InstanceNorm2d(out_channels), \n            # Activation Func\n            nn.LeakyReLU(0.2, inplace=True), \n        )\n    \n    # define behaviors (methods) -> feed forward \n    def forward(self, x):\n        return self.cnn(x)\n","metadata":{"execution":{"iopub.status.busy":"2024-07-16T20:42:05.252238Z","iopub.execute_input":"2024-07-16T20:42:05.252585Z","iopub.status.idle":"2024-07-16T20:42:05.259877Z","shell.execute_reply.started":"2024-07-16T20:42:05.252559Z","shell.execute_reply":"2024-07-16T20:42:05.258822Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Discriminator\n* Initial Convolution: The network starts with a convolution layer to process the initial input. This layer uses a stride of 2 to reduce dimensions and applies a leaky ReLU for non-linear activation.\n\n* Blocks: The discriminator sequentially applies several Block instances, each performing a convolution, normalization, and activation. These layers progressively increase the number of channels while reducing the spatial size (depending on the stride).\n\n* Final Convolution: A convolutional layer that maps the deep feature representation to a single output channel.\n* Sigmoid Activation: After the final layer, a sigmoid activation function is applied to squish the output to a [0,1] range, indicating the probability of the input being real.","metadata":{}},{"cell_type":"code","source":"class Discriminator(nn.Module):\n    def __init__(self, in_channels=3, features=[64, 128, 256, 512]):\n        super().__init__()\n        self.initial = nn.Sequential(\n            nn.Conv2d(\n                in_channels,\n                features[0],\n                kernel_size=4,\n                stride=2,\n                padding=1,\n                padding_mode=\"reflect\",\n            ),\n            nn.LeakyReLU(0.2, inplace=True),\n        )\n\n        layers = []\n        in_channels = features[0]\n        for feature in features[1:]:\n            layers.append(\n                Block(in_channels, feature, stride=1 if feature == features[-1] else 2)\n            )\n            in_channels = feature\n        layers.append(\n            nn.Conv2d(\n                in_channels,\n                1,\n                kernel_size=4,\n                stride=1,\n                padding=1,\n                padding_mode=\"reflect\",\n            )\n        )\n        # * unpack layers list into separate arguments\n        # layer in layers forms feed-forward NNs\n        self.model = nn.Sequential(*layers)\n\n    \n    # define how the input tensor x pass thru the NNs in forward:\n    # x -> first pass thru the initial layer defined above\n    # -> then thru the model layer(dynamically added based on the features list)\n    # -> sigmoid last layer to map output to 0~1 probability (fake~real)\n    def forward(self, x):\n        x = self.initial(x)\n        return torch.sigmoid(self.model(x))\n\n","metadata":{"execution":{"iopub.status.busy":"2024-07-16T20:42:07.775056Z","iopub.execute_input":"2024-07-16T20:42:07.775758Z","iopub.status.idle":"2024-07-16T20:42:07.785236Z","shell.execute_reply.started":"2024-07-16T20:42:07.775727Z","shell.execute_reply":"2024-07-16T20:42:07.784266Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Test Discriminator\n\nmodel_dis_test = Discriminator()\n\n# creates random tensor x to simulates a batch of 2 images\ntest_input = torch.randn((2, 3, 256, 256))\n\n# Forward the test input through the model\nprint(model_dis_test(test_input).shape)\n","metadata":{"execution":{"iopub.status.busy":"2024-07-16T20:42:08.973346Z","iopub.execute_input":"2024-07-16T20:42:08.973739Z","iopub.status.idle":"2024-07-16T20:42:09.101940Z","shell.execute_reply.started":"2024-07-16T20:42:08.973705Z","shell.execute_reply":"2024-07-16T20:42:09.101038Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Define CNNBlock and ResidualBlock \n* **CNNBlock** is used to build both ***downsampling*** (standard convolution -> reduces the spatial dimensions while increasing feature depth to abstract higher-level features) and ***upsampling*** (transposed convolution -> increases the spatial dimensions while refining features to reconstruct detailed and high-quality images) layers within the Generator.\n* **ResidualBlock** adds input of the block to its output, mitigating the vanishing gradient problem ->  important for training deeper networks effectively; and enabling the network to learn identity mappings efficiently -> ensuring that essential information is not lost as data passes through multiple layers of transformations.","metadata":{}},{"cell_type":"code","source":"class CNNBlock(nn.Module):\n    def __init__(self, in_channels, out_channels, down=True, use_act=True, **kwargs):\n        super().__init__()\n        \n        if down:\n            conv_layer = nn.Conv2d(in_channels, out_channels, padding_mode=\"reflect\", **kwargs)\n        else:\n            conv_layer = nn.ConvTranspose2d(in_channels, out_channels, **kwargs)\n        \n        if use_act:\n            act_layer = nn.ReLU(inplace=True)\n        else:\n            act_layer = nn.Identity()\n        \n        self.cnn = nn.Sequential(\n            conv_layer,\n            nn.InstanceNorm2d(out_channels),\n            act_layer,\n        )\n\n    def forward(self, x):\n        return self.cnn(x)\n\n\nclass ResidualBlock(nn.Module):\n    def __init__(self, channels):\n        super().__init__()\n        self.block = nn.Sequential(\n            CNNBlock(channels, channels, kernel_size=3, padding=1),\n            CNNBlock(channels, channels, use_act=False, kernel_size=3, padding=1),\n        )\n\n    def forward(self, x):\n        return x + self.block(x)\n","metadata":{"execution":{"iopub.status.busy":"2024-07-16T20:42:12.603971Z","iopub.execute_input":"2024-07-16T20:42:12.604314Z","iopub.status.idle":"2024-07-16T20:42:12.613313Z","shell.execute_reply.started":"2024-07-16T20:42:12.604287Z","shell.execute_reply":"2024-07-16T20:42:12.612410Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Generator\n* Initial Convolution: Processes the initial input to expand the feature space without changing spatial dimensions, using a 7x7 kernel size with a stride of 1 and padding set to 3. \n* Downsampling Blocks: Reduces the spatial dimensions while increasing the feature depth, using multiple CNNBlock instances each with a stride of 2. Halve dimensions and double feature channels.\n* Residual Blocks: Use multiple ResidualBlock instances to enhances feature representations without altering their spatial dimensions.\n* Upsampling Blocks: Restores the spatial dimensions to match the original input size while refining features, ussing CNNBlock instances configured for upsampling through transposed convolution. \n* Final Convolution: Reconstructs the final image from the refined feature maps, employing a large 7x7 convolutional layer to ensure detailed spatial capture. Use tanh activation function to map output to [-1,1]","metadata":{}},{"cell_type":"code","source":"class Generator(nn.Module):\n    def __init__(self, img_channels, num_features=64, num_residuals=9):\n        super().__init__()\n        self.initial = nn.Sequential(\n            nn.Conv2d(\n                img_channels,\n                num_features,\n                kernel_size=7,\n                stride=1,\n                padding=3,\n                padding_mode=\"reflect\",\n            ),\n            nn.InstanceNorm2d(num_features),\n            nn.ReLU(inplace=True),\n        )\n        self.down_blocks = nn.ModuleList(\n            [\n                CNNBlock(\n                    num_features, \n                    num_features * 2, \n                    kernel_size=3, \n                    stride=2, \n                    padding=1\n                ),\n                CNNBlock(\n                    num_features * 2,\n                    num_features * 4,\n                    kernel_size=3,\n                    stride=2,\n                    padding=1,\n                ),\n            ]\n        )\n        self.res_blocks = nn.Sequential(\n            *[ResidualBlock(num_features * 4) for _ in range(num_residuals)]\n        )\n        self.up_blocks = nn.ModuleList(\n            [\n                CNNBlock(\n                    num_features * 4,\n                    num_features * 2,\n                    down=False,\n                    kernel_size=3,\n                    stride=2,\n                    padding=1,\n                    output_padding=1,\n                ),\n                CNNBlock(\n                    num_features * 2,\n                    num_features * 1,\n                    down=False,\n                    kernel_size=3,\n                    stride=2,\n                    padding=1,\n                    output_padding=1,\n                ),\n            ]\n        )\n\n        self.last = nn.Conv2d(\n            num_features * 1,\n            img_channels,\n            kernel_size=7,\n            stride=1,\n            padding=3,\n            padding_mode=\"reflect\",\n        )\n\n    def forward(self, x):\n        x = self.initial(x)\n        for layer in self.down_blocks:\n            x = layer(x)\n        x = self.res_blocks(x)\n        for layer in self.up_blocks:\n            x = layer(x)\n        return torch.tanh(self.last(x))\n","metadata":{"execution":{"iopub.status.busy":"2024-07-16T20:42:14.009239Z","iopub.execute_input":"2024-07-16T20:42:14.009853Z","iopub.status.idle":"2024-07-16T20:42:14.021622Z","shell.execute_reply.started":"2024-07-16T20:42:14.009812Z","shell.execute_reply":"2024-07-16T20:42:14.020688Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Test Generator \n\nimg_channels = 3\nimg_size = 256\n# creates random tensor x to simulates a batch of 2 images\ntest_input = torch.randn((2, img_channels, img_size, img_size))\n\nmodel_gen_test = Generator(img_channels, 9)\n\n# Forward the test input through the model\nprint(model_gen_test(test_input).shape)","metadata":{"execution":{"iopub.status.busy":"2024-07-16T20:42:15.886171Z","iopub.execute_input":"2024-07-16T20:42:15.886809Z","iopub.status.idle":"2024-07-16T20:42:16.022196Z","shell.execute_reply.started":"2024-07-16T20:42:15.886776Z","shell.execute_reply":"2024-07-16T20:42:16.021201Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Training","metadata":{}},{"cell_type":"code","source":"# Check for available device\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nprint(f\"Using device: {device}\")","metadata":{"execution":{"iopub.status.busy":"2024-07-16T20:42:18.505971Z","iopub.execute_input":"2024-07-16T20:42:18.506810Z","iopub.status.idle":"2024-07-16T20:42:18.511521Z","shell.execute_reply.started":"2024-07-16T20:42:18.506776Z","shell.execute_reply":"2024-07-16T20:42:18.510443Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Train_Val_Split","metadata":{}},{"cell_type":"code","source":"from torch.utils.data import random_split\n\n# Define the size of the validation set\nval_size_monet = int(0.2 * len(MonetDataset))\ntrain_size_monet = len(MonetDataset) - val_size_monet\n\nval_size_photo = int(0.2 * len(PhotoDataset))\ntrain_size_photo = len(PhotoDataset) - val_size_photo\n\n\n# Split datasets into training and validation\ntrain_MonetDataset, val_MonetDataset = random_split(MonetDataset, [train_size_monet, val_size_monet])\ntrain_PhotoDataset, val_PhotoDataset = random_split(PhotoDataset, [train_size_photo, val_size_photo])\n\n# Create data loaders for training and validation datasets\ntrain_MonetLoader = DataLoader(train_MonetDataset, batch_size=8, shuffle=True, num_workers=2)\nval_MonetLoader = DataLoader(val_MonetDataset, batch_size=8, shuffle=False, num_workers=2)\n\ntrain_PhotoLoader = DataLoader(train_PhotoDataset, batch_size=8, shuffle=True, num_workers=2)\nval_PhotoLoader = DataLoader(val_PhotoDataset, batch_size=8, shuffle=False, num_workers=2)","metadata":{"execution":{"iopub.status.busy":"2024-07-16T20:42:19.893874Z","iopub.execute_input":"2024-07-16T20:42:19.894252Z","iopub.status.idle":"2024-07-16T20:42:19.902428Z","shell.execute_reply.started":"2024-07-16T20:42:19.894222Z","shell.execute_reply":"2024-07-16T20:42:19.901516Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import torch.optim as optim\nfrom torch.cuda.amp import GradScaler, autocast\nfrom itertools import product\n\n# Function to create model, optimizer, and loss functions\ndef create_model_and_optimizer(lr, betas):\n    G_Photo2Monet = Generator(img_channels=3).to(device)\n    G_Monet2Photo = Generator(img_channels=3).to(device)\n    D_Photo = Discriminator(in_channels=3).to(device)\n    D_Monet = Discriminator(in_channels=3).to(device)\n\n    optimizer_G_Photo2Monet = optim.Adam(G_Photo2Monet.parameters(), lr=lr, betas=betas)\n    optimizer_G_Monet2Photo = optim.Adam(G_Monet2Photo.parameters(), lr=lr, betas=betas)\n    optimizer_D_Photo = optim.Adam(D_Photo.parameters(), lr=lr, betas=betas)\n    optimizer_D_Monet = optim.Adam(D_Monet.parameters(), lr=lr, betas=betas)\n\n    return (G_Photo2Monet, G_Monet2Photo, D_Photo, D_Monet,\n            optimizer_G_Photo2Monet, optimizer_G_Monet2Photo,\n            optimizer_D_Photo, optimizer_D_Monet)\n\n# Loss functions\ncriterion_GAN = nn.MSELoss()\ncriterion_cycle = nn.L1Loss()\n\n# Hyperparameters to tune\nlearning_rates = [2e-4, 1e-4]\nbetas_list = [(0.5, 0.999), (0.5, 0.5)]\nlambda_cycle = 2\nnum_epochs = 10\nbatch_size = 8\n\n# Mixed precision training\nscaler = GradScaler()\n\n# Hyperparameter tuning loop\nbest_loss = float('inf')\nbest_hyperparams = None\n\nfor lr, betas in product(learning_rates, betas_list):\n    print(f\"Training with learning rate: {lr} and betas: {betas}\")\n    \n    # Initialize models and optimizers\n    (G_Photo2Monet, G_Monet2Photo, D_Photo, D_Monet,\n     optimizer_G_Photo2Monet, optimizer_G_Monet2Photo,\n     optimizer_D_Photo, optimizer_D_Monet) = create_model_and_optimizer(lr, betas)\n    \n    for epoch in range(num_epochs):\n        epoch_loss_G_Photo2Monet = 0\n        epoch_loss_G_Monet2Photo = 0\n        epoch_loss_D_Photo = 0\n        epoch_loss_D_Monet = 0\n        \n        for i, (real_photo, real_monet) in enumerate(zip(train_PhotoLoader, train_MonetLoader)):\n            real_photo = real_photo.to(device)\n            real_monet = real_monet.to(device)\n            batch_size = real_photo.size(0)\n\n            # Generate fake images\n            fake_monet = G_Photo2Monet(real_photo).detach()\n            fake_photo = G_Monet2Photo(real_monet).detach()\n\n            # Get discriminator output sizes\n            D_Photo_real = D_Photo(real_photo)\n            D_Monet_real = D_Monet(real_monet)\n            D_Photo_output_shape = D_Photo_real.shape\n            D_Monet_output_shape = D_Monet_real.shape\n\n            # Create labels with the correct shape\n            real_label_Photo = torch.ones(D_Photo_output_shape, dtype=torch.float, device=device)\n            fake_label_Photo = torch.zeros(D_Photo_output_shape, dtype=torch.float, device=device)\n            real_label_Monet = torch.ones(D_Monet_output_shape, dtype=torch.float, device=device)\n            fake_label_Monet = torch.zeros(D_Monet_output_shape, dtype=torch.float, device=device)\n\n            # ----------------------\n            # Train Discriminators\n            # ----------------------\n\n            with autocast():\n                # Discriminator D_Photo\n                D_Photo_real = D_Photo(real_photo)\n                D_Photo_fake = D_Photo(fake_photo)\n                loss_D_Photo_real = criterion_GAN(D_Photo_real, real_label_Photo)\n                loss_D_Photo_fake = criterion_GAN(D_Photo_fake, fake_label_Photo)\n                loss_D_Photo = (loss_D_Photo_real + loss_D_Photo_fake) * 0.5\n\n                # Discriminator D_Monet\n                D_Monet_real = D_Monet(real_monet)\n                D_Monet_fake = D_Monet(fake_monet)\n                loss_D_Monet_real = criterion_GAN(D_Monet_real, real_label_Monet)\n                loss_D_Monet_fake = criterion_GAN(D_Monet_fake, fake_label_Monet)\n                loss_D_Monet = (loss_D_Monet_real + loss_D_Monet_fake) * 0.5\n\n            optimizer_D_Photo.zero_grad()\n            scaler.scale(loss_D_Photo).backward()\n            scaler.step(optimizer_D_Photo)\n\n            optimizer_D_Monet.zero_grad()\n            scaler.scale(loss_D_Monet).backward()\n            scaler.step(optimizer_D_Monet)\n\n            # Accumulate discriminator loss\n            epoch_loss_D_Photo += loss_D_Photo.item()\n            epoch_loss_D_Monet += loss_D_Monet.item()\n\n            # ----------------------\n            # Train Generators\n            # ----------------------\n\n            with autocast():\n                # GAN loss\n                fake_monet = G_Photo2Monet(real_photo)\n                fake_photo = G_Monet2Photo(real_monet)\n                loss_GAN_Photo2Monet = criterion_GAN(D_Monet(fake_monet), real_label_Monet)\n                loss_GAN_Monet2Photo = criterion_GAN(D_Photo(fake_photo), real_label_Photo)\n\n                # Cycle loss\n                cycle_photo = G_Monet2Photo(fake_monet)\n                cycle_monet = G_Photo2Monet(fake_photo)\n                loss_cycle_photo = criterion_cycle(cycle_photo, real_photo)\n                loss_cycle_monet = criterion_cycle(cycle_monet, real_monet)\n\n                # Total generator loss\n                loss_G_Photo2Monet = loss_GAN_Photo2Monet + lambda_cycle * loss_cycle_photo\n                loss_G_Monet2Photo = loss_GAN_Monet2Photo + lambda_cycle * loss_cycle_monet\n\n            optimizer_G_Photo2Monet.zero_grad()\n            scaler.scale(loss_G_Photo2Monet).backward()\n            scaler.step(optimizer_G_Photo2Monet)\n\n            optimizer_G_Monet2Photo.zero_grad()\n            scaler.scale(loss_G_Monet2Photo).backward()\n            scaler.step(optimizer_G_Monet2Photo)\n\n            scaler.update()\n\n            # Accumulate generator loss\n            epoch_loss_G_Photo2Monet += loss_G_Photo2Monet.item()\n            epoch_loss_G_Monet2Photo += loss_G_Monet2Photo.item()\n\n        # Calculate average losses for the epoch\n        epoch_loss_G_Photo2Monet /= len(train_PhotoLoader)\n        epoch_loss_G_Monet2Photo /= len(train_MonetLoader)\n        epoch_loss_D_Photo /= len(train_PhotoLoader)\n        epoch_loss_D_Monet /= len(train_MonetLoader)\n\n        print(f\"Epoch [{epoch+1}/{num_epochs}] - \"\n              f\"loss_G_Photo2Monet: {epoch_loss_G_Photo2Monet:.4f}, \"\n              f\"loss_G_Monet2Photo: {epoch_loss_G_Monet2Photo:.4f}, \"\n              f\"loss_D_Photo: {epoch_loss_D_Photo:.4f}, \"\n              f\"loss_D_Monet: {epoch_loss_D_Monet:.4f}\")\n\n        # Validation and display results\n        with torch.no_grad():\n            G_Photo2Monet.eval()\n            example_photo = next(iter(val_PhotoLoader)).to(device)\n            fake_monet = G_Photo2Monet(example_photo)\n\n            # Display images\n            plt.figure(figsize=(12, 6))\n            for i in range(4):\n                plt.subplot(2, 4, i + 1)\n                plt.imshow(((example_photo[i].cpu().numpy().transpose(1, 2, 0) + 1) / 2))\n                plt.axis('off')\n                plt.subplot(2, 4, i + 5)\n                plt.imshow(((fake_monet[i].cpu().numpy().transpose(1, 2, 0) + 1) / 2))\n                plt.axis('off')\n            plt.show()\n\n            G_Photo2Monet.train()\n\n    # Check if current hyperparameters are better\n    if epoch_loss_G_Photo2Monet < best_loss:\n        best_loss = epoch_loss_G_Photo2Monet\n        best_hyperparams = (lr, betas)\n\n    print(f\"Best hyperparameters: learning rate = {best_hyperparams[0]}, betas = {best_hyperparams[1]}\")\n\n# Save final model\nfinal_model_path = 'final_model.pth'\ntorch.save({\n    'G_Photo2Monet': G_Photo2Monet.state_dict(),\n    'G_Monet2Photo': G_Monet2Photo.state_dict(),\n    'D_Photo': D_Photo.state_dict(),\n    'D_Monet': D_Monet.state_dict(),\n    'optimizer_G_Photo2Monet': optimizer_G_Photo2Monet.state_dict(),\n    'optimizer_G_Monet2Photo': optimizer_G_Monet2Photo.state_dict(),\n    'optimizer_D_Photo': optimizer_D_Photo.state_dict(),\n    'optimizer_D_Monet': optimizer_D_Monet.state_dict(),\n}, final_model_path)\n","metadata":{"execution":{"iopub.status.busy":"2024-07-16T20:43:34.381072Z","iopub.execute_input":"2024-07-16T20:43:34.381443Z","iopub.status.idle":"2024-07-16T20:49:20.458636Z","shell.execute_reply.started":"2024-07-16T20:43:34.381415Z","shell.execute_reply":"2024-07-16T20:49:20.457469Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"\n# Load the final model from the previous training\ncheckpoint_path = 'final_model.pth'\ncheckpoint = torch.load(checkpoint_path)\n\n# Initialize models and optimizers with best hyperparameters\nbest_lr = 0.0001\nbest_betas = (0.5, 0.5)\n\nG_Photo2Monet = Generator(img_channels=3).to(device)\nG_Monet2Photo = Generator(img_channels=3).to(device)\nD_Photo = Discriminator(in_channels=3).to(device)\nD_Monet = Discriminator(in_channels=3).to(device)\n\nG_Photo2Monet.load_state_dict(checkpoint['G_Photo2Monet'])\nG_Monet2Photo.load_state_dict(checkpoint['G_Monet2Photo'])\nD_Photo.load_state_dict(checkpoint['D_Photo'])\nD_Monet.load_state_dict(checkpoint['D_Monet'])\n\noptimizer_G_Photo2Monet = optim.Adam(G_Photo2Monet.parameters(), lr=best_lr, betas=best_betas)\noptimizer_G_Monet2Photo = optim.Adam(G_Monet2Photo.parameters(), lr=best_lr, betas=best_betas)\noptimizer_D_Photo = optim.Adam(D_Photo.parameters(), lr=best_lr, betas=best_betas)\noptimizer_D_Monet = optim.Adam(D_Monet.parameters(), lr=best_lr, betas=best_betas)\n\noptimizer_G_Photo2Monet.load_state_dict(checkpoint['optimizer_G_Photo2Monet'])\noptimizer_G_Monet2Photo.load_state_dict(checkpoint['optimizer_G_Monet2Photo'])\noptimizer_D_Photo.load_state_dict(checkpoint['optimizer_D_Photo'])\noptimizer_D_Monet.load_state_dict(checkpoint['optimizer_D_Monet'])\n\n# Loss functions\ncriterion_GAN = nn.MSELoss()\ncriterion_cycle = nn.L1Loss()\n\n# Mixed precision training\nscaler = GradScaler()\n\n# Continue training for an additional 20 epochs\nnum_epochs = 20\n\nfor epoch in range(num_epochs):\n    epoch_loss_G_Photo2Monet = 0\n    epoch_loss_G_Monet2Photo = 0\n    epoch_loss_D_Photo = 0\n    epoch_loss_D_Monet = 0\n    \n    for i, (real_photo, real_monet) in enumerate(zip(train_PhotoLoader, train_MonetLoader)):\n        real_photo = real_photo.to(device)\n        real_monet = real_monet.to(device)\n        batch_size = real_photo.size(0)\n\n        # Generate fake images\n        fake_monet = G_Photo2Monet(real_photo).detach()\n        fake_photo = G_Monet2Photo(real_monet).detach()\n\n        # Get discriminator output sizes\n        D_Photo_real = D_Photo(real_photo)\n        D_Monet_real = D_Monet(real_monet)\n        D_Photo_output_shape = D_Photo_real.shape\n        D_Monet_output_shape = D_Monet_real.shape\n\n        # Create labels with the correct shape\n        real_label_Photo = torch.ones(D_Photo_output_shape, dtype=torch.float, device=device)\n        fake_label_Photo = torch.zeros(D_Photo_output_shape, dtype=torch.float, device=device)\n        real_label_Monet = torch.ones(D_Monet_output_shape, dtype=torch.float, device=device)\n        fake_label_Monet = torch.zeros(D_Monet_output_shape, dtype=torch.float, device=device)\n\n        # ----------------------\n        # Train Discriminators\n        # ----------------------\n\n        with autocast():\n            # Discriminator D_Photo\n            D_Photo_real = D_Photo(real_photo)\n            D_Photo_fake = D_Photo(fake_photo)\n            loss_D_Photo_real = criterion_GAN(D_Photo_real, real_label_Photo)\n            loss_D_Photo_fake = criterion_GAN(D_Photo_fake, fake_label_Photo)\n            loss_D_Photo = (loss_D_Photo_real + loss_D_Photo_fake) * 0.5\n\n            # Discriminator D_Monet\n            D_Monet_real = D_Monet(real_monet)\n            D_Monet_fake = D_Monet(fake_monet)\n            loss_D_Monet_real = criterion_GAN(D_Monet_real, real_label_Monet)\n            loss_D_Monet_fake = criterion_GAN(D_Monet_fake, fake_label_Monet)\n            loss_D_Monet = (loss_D_Monet_real + loss_D_Monet_fake) * 0.5\n\n        optimizer_D_Photo.zero_grad()\n        scaler.scale(loss_D_Photo).backward()\n        scaler.step(optimizer_D_Photo)\n\n        optimizer_D_Monet.zero_grad()\n        scaler.scale(loss_D_Monet).backward()\n        scaler.step(optimizer_D_Monet)\n\n        # Accumulate discriminator loss\n        epoch_loss_D_Photo += loss_D_Photo.item()\n        epoch_loss_D_Monet += loss_D_Monet.item()\n\n        # ----------------------\n        # Train Generators\n        # ----------------------\n\n        with autocast():\n            # GAN loss\n            fake_monet = G_Photo2Monet(real_photo)\n            fake_photo = G_Monet2Photo(real_monet)\n            loss_GAN_Photo2Monet = criterion_GAN(D_Monet(fake_monet), real_label_Monet)\n            loss_GAN_Monet2Photo = criterion_GAN(D_Photo(fake_photo), real_label_Photo)\n\n            # Cycle loss\n            cycle_photo = G_Monet2Photo(fake_monet)\n            cycle_monet = G_Photo2Monet(fake_photo)\n            loss_cycle_photo = criterion_cycle(cycle_photo, real_photo)\n            loss_cycle_monet = criterion_cycle(cycle_monet, real_monet)\n\n            # Total generator loss\n            loss_G_Photo2Monet = loss_GAN_Photo2Monet + lambda_cycle * loss_cycle_photo\n            loss_G_Monet2Photo = loss_GAN_Monet2Photo + lambda_cycle * loss_cycle_monet\n\n        optimizer_G_Photo2Monet.zero_grad()\n        scaler.scale(loss_G_Photo2Monet).backward()\n        scaler.step(optimizer_G_Photo2Monet)\n\n        optimizer_G_Monet2Photo.zero_grad()\n        scaler.scale(loss_G_Monet2Photo).backward()\n        scaler.step(optimizer_G_Monet2Photo)\n\n        scaler.update()\n\n        # Accumulate generator loss\n        epoch_loss_G_Photo2Monet += loss_G_Photo2Monet.item()\n        epoch_loss_G_Monet2Photo += loss_G_Monet2Photo.item()\n\n    # Calculate average losses for the epoch\n    epoch_loss_G_Photo2Monet /= len(train_PhotoLoader)\n    epoch_loss_G_Monet2Photo /= len(train_MonetLoader)\n    epoch_loss_D_Photo /= len(train_PhotoLoader)\n    epoch_loss_D_Monet /= len(train_MonetLoader)\n\n    print(f\"Epoch [{epoch+1}/{num_epochs}] - \"\n          f\"loss_G_Photo2Monet: {epoch_loss_G_Photo2Monet:.4f}, \"\n          f\"loss_G_Monet2Photo: {epoch_loss_G_Monet2Photo:.4f}, \"\n          f\"loss_D_Photo: {epoch_loss_D_Photo:.4f}, \"\n          f\"loss_D_Monet: {epoch_loss_D_Monet:.4f}\")\n\n    # Validation and display results\n    with torch.no_grad():\n        G_Photo2Monet.eval()\n        example_photo = next(iter(val_PhotoLoader)).to(device)\n        fake_monet = G_Photo2Monet(example_photo)\n\n        # Display images\n        plt.figure(figsize=(12, 6))\n        for i in range(4):\n            plt.subplot(2, 4, i + 1)\n            plt.imshow(((example_photo[i].cpu().numpy().transpose(1, 2, 0) + 1) / 2))\n            plt.axis('off')\n            plt.subplot(2, 4, i + 5)\n            plt.imshow(((fake_monet[i].cpu().numpy().transpose(1, 2, 0) + 1) / 2))\n            plt.axis('off')\n        plt.show()\n\n        G_Photo2Monet.train()\n\n# Save final model\nfinal_model_v2_path = 'final_model_v2.pth'\ntorch.save({\n    'G_Photo2Monet': G_Photo2Monet.state_dict(),\n    'G_Monet2Photo': G_Monet2Photo.state_dict(),\n    'D_Photo': D_Photo.state_dict(),\n    'D_Monet': D_Monet.state_dict(),\n    'optimizer_G_Photo2Monet': optimizer_G_Photo2Monet.state_dict(),\n    'optimizer_G_Monet2Photo': optimizer_G_Monet2Photo.state_dict(),\n    'optimizer_D_Photo': optimizer_D_Photo.state_dict(),\n    'optimizer_D_Monet': optimizer_D_Monet.state_dict(),\n}, final_model_v2_path)\n","metadata":{"execution":{"iopub.status.busy":"2024-07-16T21:09:56.351857Z","iopub.execute_input":"2024-07-16T21:09:56.352514Z","iopub.status.idle":"2024-07-16T21:12:49.573480Z","shell.execute_reply.started":"2024-07-16T21:09:56.352475Z","shell.execute_reply":"2024-07-16T21:12:49.572601Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Adjust the lambda_cycle and continue training \n\nThe lambda_cycle determines how much weight the cycle consistency loss has in the model training process.","metadata":{}},{"cell_type":"code","source":"# Load the final model from the previous training\ncheckpoint_path = 'final_model_v2.pth'\ncheckpoint = torch.load(checkpoint_path)\n\n# Initialize models and optimizers with best hyperparameters\nbest_lr = 0.0001\nbest_betas = (0.5, 0.5)\n\nG_Photo2Monet = Generator(img_channels=3).to(device)\nG_Monet2Photo = Generator(img_channels=3).to(device)\nD_Photo = Discriminator(in_channels=3).to(device)\nD_Monet = Discriminator(in_channels=3).to(device)\n\nG_Photo2Monet.load_state_dict(checkpoint['G_Photo2Monet'])\nG_Monet2Photo.load_state_dict(checkpoint['G_Monet2Photo'])\nD_Photo.load_state_dict(checkpoint['D_Photo'])\nD_Monet.load_state_dict(checkpoint['D_Monet'])\n\noptimizer_G_Photo2Monet = optim.Adam(G_Photo2Monet.parameters(), lr=best_lr, betas=best_betas)\noptimizer_G_Monet2Photo = optim.Adam(G_Monet2Photo.parameters(), lr=best_lr, betas=best_betas)\noptimizer_D_Photo = optim.Adam(D_Photo.parameters(), lr=best_lr, betas=best_betas)\noptimizer_D_Monet = optim.Adam(D_Monet.parameters(), lr=best_lr, betas=best_betas)\n\noptimizer_G_Photo2Monet.load_state_dict(checkpoint['optimizer_G_Photo2Monet'])\noptimizer_G_Monet2Photo.load_state_dict(checkpoint['optimizer_G_Monet2Photo'])\noptimizer_D_Photo.load_state_dict(checkpoint['optimizer_D_Photo'])\noptimizer_D_Monet.load_state_dict(checkpoint['optimizer_D_Monet'])\n\n# Loss functions\ncriterion_GAN = nn.MSELoss()\ncriterion_cycle = nn.L1Loss()\n\n# Mixed precision training\nscaler = GradScaler()\n\n# Adjust the lambda_cycle\nlambda_cycle = 5\n\n# Continue training for an additional 20 epochs\nnum_epochs = 20\n\nfor epoch in range(num_epochs):\n    epoch_loss_G_Photo2Monet = 0\n    epoch_loss_G_Monet2Photo = 0\n    epoch_loss_D_Photo = 0\n    epoch_loss_D_Monet = 0\n    \n    for i, (real_photo, real_monet) in enumerate(zip(train_PhotoLoader, train_MonetLoader)):\n        real_photo = real_photo.to(device)\n        real_monet = real_monet.to(device)\n        batch_size = real_photo.size(0)\n\n        # Generate fake images\n        fake_monet = G_Photo2Monet(real_photo).detach()\n        fake_photo = G_Monet2Photo(real_monet).detach()\n\n        # Get discriminator output sizes\n        D_Photo_real = D_Photo(real_photo)\n        D_Monet_real = D_Monet(real_monet)\n        D_Photo_output_shape = D_Photo_real.shape\n        D_Monet_output_shape = D_Monet_real.shape\n\n        # Create labels with the correct shape\n        real_label_Photo = torch.ones(D_Photo_output_shape, dtype=torch.float, device=device)\n        fake_label_Photo = torch.zeros(D_Photo_output_shape, dtype=torch.float, device=device)\n        real_label_Monet = torch.ones(D_Monet_output_shape, dtype=torch.float, device=device)\n        fake_label_Monet = torch.zeros(D_Monet_output_shape, dtype=torch.float, device=device)\n\n        # ----------------------\n        # Train Discriminators\n        # ----------------------\n\n        with autocast():\n            # Discriminator D_Photo\n            D_Photo_real = D_Photo(real_photo)\n            D_Photo_fake = D_Photo(fake_photo)\n            loss_D_Photo_real = criterion_GAN(D_Photo_real, real_label_Photo)\n            loss_D_Photo_fake = criterion_GAN(D_Photo_fake, fake_label_Photo)\n            loss_D_Photo = (loss_D_Photo_real + loss_D_Photo_fake) * 0.5\n\n            # Discriminator D_Monet\n            D_Monet_real = D_Monet(real_monet)\n            D_Monet_fake = D_Monet(fake_monet)\n            loss_D_Monet_real = criterion_GAN(D_Monet_real, real_label_Monet)\n            loss_D_Monet_fake = criterion_GAN(D_Monet_fake, fake_label_Monet)\n            loss_D_Monet = (loss_D_Monet_real + loss_D_Monet_fake) * 0.5\n\n        optimizer_D_Photo.zero_grad()\n        scaler.scale(loss_D_Photo).backward()\n        scaler.step(optimizer_D_Photo)\n\n        optimizer_D_Monet.zero_grad()\n        scaler.scale(loss_D_Monet).backward()\n        scaler.step(optimizer_D_Monet)\n\n        # Accumulate discriminator loss\n        epoch_loss_D_Photo += loss_D_Photo.item()\n        epoch_loss_D_Monet += loss_D_Monet.item()\n\n        # ----------------------\n        # Train Generators\n        # ----------------------\n\n        with autocast():\n            # GAN loss\n            fake_monet = G_Photo2Monet(real_photo)\n            fake_photo = G_Monet2Photo(real_monet)\n            loss_GAN_Photo2Monet = criterion_GAN(D_Monet(fake_monet), real_label_Monet)\n            loss_GAN_Monet2Photo = criterion_GAN(D_Photo(fake_photo), real_label_Photo)\n\n            # Cycle loss\n            cycle_photo = G_Monet2Photo(fake_monet)\n            cycle_monet = G_Photo2Monet(fake_photo)\n            loss_cycle_photo = criterion_cycle(cycle_photo, real_photo)\n            loss_cycle_monet = criterion_cycle(cycle_monet, real_monet)\n\n            # Total generator loss\n            loss_G_Photo2Monet = loss_GAN_Photo2Monet + lambda_cycle * loss_cycle_photo\n            loss_G_Monet2Photo = loss_GAN_Monet2Photo + lambda_cycle * loss_cycle_monet\n\n        optimizer_G_Photo2Monet.zero_grad()\n        scaler.scale(loss_G_Photo2Monet).backward()\n        scaler.step(optimizer_G_Photo2Monet)\n\n        optimizer_G_Monet2Photo.zero_grad()\n        scaler.scale(loss_G_Monet2Photo).backward()\n        scaler.step(optimizer_G_Monet2Photo)\n\n        scaler.update()\n\n        # Accumulate generator loss\n        epoch_loss_G_Photo2Monet += loss_G_Photo2Monet.item()\n        epoch_loss_G_Monet2Photo += loss_G_Monet2Photo.item()\n\n    # Calculate average losses for the epoch\n    epoch_loss_G_Photo2Monet /= len(train_PhotoLoader)\n    epoch_loss_G_Monet2Photo /= len(train_MonetLoader)\n    epoch_loss_D_Photo /= len(train_PhotoLoader)\n    epoch_loss_D_Monet /= len(train_MonetLoader)\n\n    print(f\"Epoch [{epoch+1}/{num_epochs}] - \"\n          f\"loss_G_Photo2Monet: {epoch_loss_G_Photo2Monet:.4f}, \"\n          f\"loss_G_Monet2Photo: {epoch_loss_G_Monet2Photo:.4f}, \"\n          f\"loss_D_Photo: {epoch_loss_D_Photo:.4f}, \"\n          f\"loss_D_Monet: {epoch_loss_D_Monet:.4f}\")\n\n    # Validation and display results\n    with torch.no_grad():\n        G_Photo2Monet.eval()\n        example_photo = next(iter(val_PhotoLoader)).to(device)\n        fake_monet = G_Photo2Monet(example_photo)\n\n        # Display images\n        plt.figure(figsize=(12, 6))\n        for i in range(4):\n            plt.subplot(2, 4, i + 1)\n            plt.imshow(((example_photo[i].cpu().numpy().transpose(1, 2, 0) + 1) / 2))\n            plt.axis('off')\n            plt.subplot(2, 4, i + 5)\n            plt.imshow(((fake_monet[i].cpu().numpy().transpose(1, 2, 0) + 1) / 2))\n            plt.axis('off')\n        plt.show()\n\n        G_Photo2Monet.train()\n\n# Save final model\nfinal_model_v3_path = 'final_model_v3.pth'\ntorch.save({\n    'G_Photo2Monet': G_Photo2Monet.state_dict(),\n    'G_Monet2Photo': G_Monet2Photo.state_dict(),\n    'D_Photo': D_Photo.state_dict(),\n    'D_Monet': D_Monet.state_dict(),\n    'optimizer_G_Photo2Monet': optimizer_G_Photo2Monet.state_dict(),\n    'optimizer_G_Monet2Photo': optimizer_G_Monet2Photo.state_dict(),\n    'optimizer_D_Photo': optimizer_D_Photo.state_dict(),\n    'optimizer_D_Monet': optimizer_D_Monet.state_dict(),\n}, final_model_v3_path)\n","metadata":{"execution":{"iopub.status.busy":"2024-07-16T21:27:47.492713Z","iopub.execute_input":"2024-07-16T21:27:47.493367Z","iopub.status.idle":"2024-07-16T21:30:40.573681Z","shell.execute_reply.started":"2024-07-16T21:27:47.493334Z","shell.execute_reply":"2024-07-16T21:30:40.572602Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Continue to train another 20 epochs","metadata":{}},{"cell_type":"code","source":"# Load the final model from the previous training\ncheckpoint_path = 'final_model_v3.pth'\ncheckpoint = torch.load(checkpoint_path)\n\n# Initialize models and optimizers with best hyperparameters\nbest_lr = 0.0001\nbest_betas = (0.5, 0.5)\n\nG_Photo2Monet = Generator(img_channels=3).to(device)\nG_Monet2Photo = Generator(img_channels=3).to(device)\nD_Photo = Discriminator(in_channels=3).to(device)\nD_Monet = Discriminator(in_channels=3).to(device)\n\nG_Photo2Monet.load_state_dict(checkpoint['G_Photo2Monet'])\nG_Monet2Photo.load_state_dict(checkpoint['G_Monet2Photo'])\nD_Photo.load_state_dict(checkpoint['D_Photo'])\nD_Monet.load_state_dict(checkpoint['D_Monet'])\n\noptimizer_G_Photo2Monet = optim.Adam(G_Photo2Monet.parameters(), lr=best_lr, betas=best_betas)\noptimizer_G_Monet2Photo = optim.Adam(G_Monet2Photo.parameters(), lr=best_lr, betas=best_betas)\noptimizer_D_Photo = optim.Adam(D_Photo.parameters(), lr=best_lr, betas=best_betas)\noptimizer_D_Monet = optim.Adam(D_Monet.parameters(), lr=best_lr, betas=best_betas)\n\noptimizer_G_Photo2Monet.load_state_dict(checkpoint['optimizer_G_Photo2Monet'])\noptimizer_G_Monet2Photo.load_state_dict(checkpoint['optimizer_G_Monet2Photo'])\noptimizer_D_Photo.load_state_dict(checkpoint['optimizer_D_Photo'])\noptimizer_D_Monet.load_state_dict(checkpoint['optimizer_D_Monet'])\n\n# Loss functions\ncriterion_GAN = nn.MSELoss()\ncriterion_cycle = nn.L1Loss()\n\n# Mixed precision training\nscaler = GradScaler()\n\n# Adjust the lambda_cycle\nlambda_cycle = 5\n\n# Continue training for an additional 20 epochs\nnum_epochs = 20\n\nfor epoch in range(num_epochs):\n    epoch_loss_G_Photo2Monet = 0\n    epoch_loss_G_Monet2Photo = 0\n    epoch_loss_D_Photo = 0\n    epoch_loss_D_Monet = 0\n    \n    for i, (real_photo, real_monet) in enumerate(zip(train_PhotoLoader, train_MonetLoader)):\n        real_photo = real_photo.to(device)\n        real_monet = real_monet.to(device)\n        batch_size = real_photo.size(0)\n\n        # Generate fake images\n        fake_monet = G_Photo2Monet(real_photo).detach()\n        fake_photo = G_Monet2Photo(real_monet).detach()\n\n        # Get discriminator output sizes\n        D_Photo_real = D_Photo(real_photo)\n        D_Monet_real = D_Monet(real_monet)\n        D_Photo_output_shape = D_Photo_real.shape\n        D_Monet_output_shape = D_Monet_real.shape\n\n        # Create labels with the correct shape\n        real_label_Photo = torch.ones(D_Photo_output_shape, dtype=torch.float, device=device)\n        fake_label_Photo = torch.zeros(D_Photo_output_shape, dtype=torch.float, device=device)\n        real_label_Monet = torch.ones(D_Monet_output_shape, dtype=torch.float, device=device)\n        fake_label_Monet = torch.zeros(D_Monet_output_shape, dtype=torch.float, device=device)\n\n        # ----------------------\n        # Train Discriminators\n        # ----------------------\n\n        with autocast():\n            # Discriminator D_Photo\n            D_Photo_real = D_Photo(real_photo)\n            D_Photo_fake = D_Photo(fake_photo)\n            loss_D_Photo_real = criterion_GAN(D_Photo_real, real_label_Photo)\n            loss_D_Photo_fake = criterion_GAN(D_Photo_fake, fake_label_Photo)\n            loss_D_Photo = (loss_D_Photo_real + loss_D_Photo_fake) * 0.5\n\n            # Discriminator D_Monet\n            D_Monet_real = D_Monet(real_monet)\n            D_Monet_fake = D_Monet(fake_monet)\n            loss_D_Monet_real = criterion_GAN(D_Monet_real, real_label_Monet)\n            loss_D_Monet_fake = criterion_GAN(D_Monet_fake, fake_label_Monet)\n            loss_D_Monet = (loss_D_Monet_real + loss_D_Monet_fake) * 0.5\n\n        optimizer_D_Photo.zero_grad()\n        scaler.scale(loss_D_Photo).backward()\n        scaler.step(optimizer_D_Photo)\n\n        optimizer_D_Monet.zero_grad()\n        scaler.scale(loss_D_Monet).backward()\n        scaler.step(optimizer_D_Monet)\n\n        # Accumulate discriminator loss\n        epoch_loss_D_Photo += loss_D_Photo.item()\n        epoch_loss_D_Monet += loss_D_Monet.item()\n\n        # ----------------------\n        # Train Generators\n        # ----------------------\n\n        with autocast():\n            # GAN loss\n            fake_monet = G_Photo2Monet(real_photo)\n            fake_photo = G_Monet2Photo(real_monet)\n            loss_GAN_Photo2Monet = criterion_GAN(D_Monet(fake_monet), real_label_Monet)\n            loss_GAN_Monet2Photo = criterion_GAN(D_Photo(fake_photo), real_label_Photo)\n\n            # Cycle loss\n            cycle_photo = G_Monet2Photo(fake_monet)\n            cycle_monet = G_Photo2Monet(fake_photo)\n            loss_cycle_photo = criterion_cycle(cycle_photo, real_photo)\n            loss_cycle_monet = criterion_cycle(cycle_monet, real_monet)\n\n            # Total generator loss\n            loss_G_Photo2Monet = loss_GAN_Photo2Monet + lambda_cycle * loss_cycle_photo\n            loss_G_Monet2Photo = loss_GAN_Monet2Photo + lambda_cycle * loss_cycle_monet\n\n        optimizer_G_Photo2Monet.zero_grad()\n        scaler.scale(loss_G_Photo2Monet).backward()\n        scaler.step(optimizer_G_Photo2Monet)\n\n        optimizer_G_Monet2Photo.zero_grad()\n        scaler.scale(loss_G_Monet2Photo).backward()\n        scaler.step(optimizer_G_Monet2Photo)\n\n        scaler.update()\n\n        # Accumulate generator loss\n        epoch_loss_G_Photo2Monet += loss_G_Photo2Monet.item()\n        epoch_loss_G_Monet2Photo += loss_G_Monet2Photo.item()\n\n    # Calculate average losses for the epoch\n    epoch_loss_G_Photo2Monet /= len(train_PhotoLoader)\n    epoch_loss_G_Monet2Photo /= len(train_MonetLoader)\n    epoch_loss_D_Photo /= len(train_PhotoLoader)\n    epoch_loss_D_Monet /= len(train_MonetLoader)\n\n    print(f\"Epoch [{epoch+1}/{num_epochs}] - \"\n          f\"loss_G_Photo2Monet: {epoch_loss_G_Photo2Monet:.4f}, \"\n          f\"loss_G_Monet2Photo: {epoch_loss_G_Monet2Photo:.4f}, \"\n          f\"loss_D_Photo: {epoch_loss_D_Photo:.4f}, \"\n          f\"loss_D_Monet: {epoch_loss_D_Monet:.4f}\")\n\n    # Validation and display results\n    with torch.no_grad():\n        G_Photo2Monet.eval()\n        example_photo = next(iter(val_PhotoLoader)).to(device)\n        fake_monet = G_Photo2Monet(example_photo)\n\n        # Display images\n        plt.figure(figsize=(12, 6))\n        for i in range(4):\n            plt.subplot(2, 4, i + 1)\n            plt.imshow(((example_photo[i].cpu().numpy().transpose(1, 2, 0) + 1) / 2))\n            plt.axis('off')\n            plt.subplot(2, 4, i + 5)\n            plt.imshow(((fake_monet[i].cpu().numpy().transpose(1, 2, 0) + 1) / 2))\n            plt.axis('off')\n        plt.show()\n\n        G_Photo2Monet.train()\n\n# Save final model\nfinal_model_v4_path = 'final_model_v4.pth'\ntorch.save({\n    'G_Photo2Monet': G_Photo2Monet.state_dict(),\n    'G_Monet2Photo': G_Monet2Photo.state_dict(),\n    'D_Photo': D_Photo.state_dict(),\n    'D_Monet': D_Monet.state_dict(),\n    'optimizer_G_Photo2Monet': optimizer_G_Photo2Monet.state_dict(),\n    'optimizer_G_Monet2Photo': optimizer_G_Monet2Photo.state_dict(),\n    'optimizer_D_Photo': optimizer_D_Photo.state_dict(),\n    'optimizer_D_Monet': optimizer_D_Monet.state_dict(),\n}, final_model_v4_path)\n","metadata":{"execution":{"iopub.status.busy":"2024-07-16T21:34:41.211208Z","iopub.execute_input":"2024-07-16T21:34:41.211585Z","iopub.status.idle":"2024-07-16T21:37:34.123848Z","shell.execute_reply.started":"2024-07-16T21:34:41.211550Z","shell.execute_reply":"2024-07-16T21:37:34.122942Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Submission","metadata":{}},{"cell_type":"code","source":"import shutil\nfrom zipfile import ZipFile\n\n# Load the trained model\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nmodel_path = '/kaggle/working/final_model_v4.pth'\ncheckpoint = torch.load(model_path)\n\nG_Photo2Monet = Generator(img_channels=3).to(device)\nG_Photo2Monet.load_state_dict(checkpoint['G_Photo2Monet'])\nG_Photo2Monet.eval()\n\n# Preprocess the test dataset\nclass TestDataset(Dataset):\n    def __init__(self, data_path, transform=None):\n        self.data_path = data_path\n        self.transform = transform\n        self.image_paths = [os.path.join(data_path, file) for file in os.listdir(data_path) if file.endswith('.jpg')]\n        \n    def __len__(self):\n        return len(self.image_paths)\n    \n    def __getitem__(self, idx):\n        img_path = self.image_paths[idx]\n        image = Image.open(img_path).convert('RGB')\n        if self.transform:\n            image = self.transform(image)\n        return image\n\n\ntest_transforms = torchvision.transforms.Compose([\n    torchvision.transforms.Resize((256,256)),\n    torchvision.transforms.ToTensor(),\n    torchvision.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))\n])\n\n\ntest_data_path = '/kaggle/input/gan-getting-started/photo_jpg'  \ntest_dataset = TestDataset(data_path=test_data_path, transform=test_transforms)\ntest_loader = DataLoader(test_dataset, batch_size=1, shuffle=False)\n\n# Create a directory to save the generated images\nos.makedirs('../tmp/images', exist_ok=True)\n\n# Generate Monet-style images and save them\nwith torch.no_grad():\n    for i, img in enumerate(test_loader):\n        img = img.to(device)\n        fake_monet = G_Photo2Monet(img)[0].cpu().permute(1, 2, 0).numpy() * 0.5 + 0.5\n        im = Image.fromarray((fake_monet * 255).astype(np.uint8))\n        im.save(f\"../tmp/images/{i}.jpg\")\n\n# Zip the generated images\noutput_zip_path = \"/kaggle/working/images.zip\"\nwith ZipFile(output_zip_path, 'w') as zipf:\n    for root, _, files in os.walk('../tmp/images'):\n        for file in files:\n            zipf.write(os.path.join(root, file), os.path.relpath(os.path.join(root, file), '../tmp/images'))\n\nprint(f\"Submission file created at {output_zip_path}\")\n","metadata":{"execution":{"iopub.status.busy":"2024-07-16T22:52:42.214198Z","iopub.execute_input":"2024-07-16T22:52:42.214873Z","iopub.status.idle":"2024-07-16T22:56:45.859916Z","shell.execute_reply.started":"2024-07-16T22:52:42.214838Z","shell.execute_reply":"2024-07-16T22:56:45.858968Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}