{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.14","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":30786,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nimport numpy as np\nimport matplotlib.pyplot as plt\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torchvision\nfrom torch.utils.data import Dataset, DataLoader, random_split\nfrom torchvision.utils import make_grid\nfrom PIL import Image\n\nfrom tqdm import tqdm\nimport itertools\nimport shutil\n\nfrom pathlib import Path\nfrom torch.cuda.amp import autocast, GradScaler","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-18T08:31:54.372091Z","iopub.execute_input":"2025-01-18T08:31:54.372369Z","iopub.status.idle":"2025-01-18T08:31:58.809038Z","shell.execute_reply.started":"2025-01-18T08:31:54.372342Z","shell.execute_reply":"2025-01-18T08:31:58.808104Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"data_root = '/kaggle/input/gan-getting-started/'\nprint(data_root)\nprint(os.listdir(data_root))\nmonet_path = 'monet_jpg'\nphoto_path = 'photo_jpg'\nBATCH_SIZE = 16\nLAMBDA_COEF = 10\nNB_EPOCHS = 20","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-18T08:31:58.810152Z","iopub.execute_input":"2025-01-18T08:31:58.810537Z","iopub.status.idle":"2025-01-18T08:31:58.818385Z","shell.execute_reply.started":"2025-01-18T08:31:58.810502Z","shell.execute_reply":"2025-01-18T08:31:58.817532Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class MonetPhoto(Dataset):\n    def __init__(self, data_root, monet_path, photo_path, transform=None):\n        \"\"\"\n        Initialize the MonetPhoto dataset with paths to Monet and photo images and optional transformations.\n        \n        Args:\n            data_root (str): Root directory where the data is stored.\n            monet_path (str): Subdirectory within data_root containing Monet-style images.\n            photo_path (str): Subdirectory within data_root containing photo images.\n            transform (callable, optional): Optional transform to be applied on a sample.\n        \"\"\"\n        self.data_root = data_root  # Set the root directory for data\n        self.monet_path = monet_path  # Set the subdirectory for Monet images\n        self.monet_images = os.listdir(os.path.join(data_root, monet_path))  # List all Monet image filenames\n        self.photo_path = photo_path  # Set the subdirectory for photo images\n        self.photo_images = os.listdir(os.path.join(data_root, photo_path))  # List all photo image filenames\n        self.transform = transform  # Set the transformation to be applied on images\n\n    def __len__(self):\n        \"\"\"\n        Return the total number of samples in the dataset.\n        The length is the maximum of the number of Monet and photo images to ensure pairing.\n        \n        Returns:\n            int: Total number of samples.\n        \"\"\"\n        return max(len(self.monet_images), len(self.photo_images))  # Ensure all images are paired\n\n    def __getitem__(self, idx):\n        \"\"\"\n        Retrieve a sample from the dataset at the given index.\n        It pairs a Monet image with a photo image. If indices exceed the number of available images,\n        it wraps around using modulo operation.\n        \n        Args:\n            idx (int): Index of the sample to retrieve.\n        \n        Returns:\n            tuple: A pair of Monet and photo images after applying transformations.\n        \"\"\"\n        # Open the Monet image at the given index, wrapping around if idx exceeds length\n        monet_image = Image.open(os.path.join(self.data_root, self.monet_path, self.monet_images[idx % len(self.monet_images)]))\n        # Open the photo image at the given index, wrapping around if idx exceeds length\n        photo_image = Image.open(os.path.join(self.data_root, self.photo_path, self.photo_images[idx % len(self.photo_images)]))\n        \n        # Apply transformations if any are provided\n        if self.transform:\n            monet_image = self.transform(monet_image)  # Transform the Monet image\n            photo_image = self.transform(photo_image)  # Transform the photo image\n        \n        return monet_image, photo_image  # Return the transformed image pair\n\n# Define a sequence of transformations to apply to the images\ntransforms = torchvision.transforms.Compose([\n    torchvision.transforms.Resize((256, 256)),  # Resize images to 256x256 pixelst\n    torchvision.transforms.RandomVerticalFlip(p=0.2),\n    torchvision.transforms.ToTensor(),  # Convert PIL images to PyTorch tensors\n    torchvision.transforms.Normalize(0.5, 0.5)  # Normalize tensor values to [-1, 1]\n])\n\n# Create an instance of the MonetPhoto dataset with the specified paths and transformations\ndataset = MonetPhoto(data_root, monet_path, photo_path, transform=transforms)\n\n# Get the total number of samples in the dataset\nlen(dataset)\n\n# Calculate the number of samples for training (80%) and testing (20%) splits\ntrain_size, test_size = int(len(dataset) * 0.8), len(dataset) - int(len(dataset) * 0.8)\n# Split the dataset into training and testing subsets\ntraining_data, testing_data = random_split(dataset, [train_size, test_size])\n\n# Get the number of samples in the training subset\nprint(len(training_data))\n\n# Create a DataLoader for the training data with a batch size of 16, shuffling enabled, and using 2 worker processes\ntrain_dataloader = DataLoader(training_data.dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=2)\n# Create a DataLoader for the testing data with a batch size of 16, shuffling enabled, and using 2 worker processes\ntest_dataloader = DataLoader(testing_data.dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=2)\n\n# Retrieve the first batch of Monet and photo images from the training DataLoader\nmones, photos = next(iter(train_dataloader))\n\n# Plot the first Monet image in the batch\nplt.imshow(mones[0].permute(1, 2, 0).numpy() * 0.5 + 0.5)  # Convert tensor to numpy array and denormalize for displayImage.open(os.path.join(self.data_root, self.photo_path, \n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-18T08:31:58.820299Z","iopub.execute_input":"2025-01-18T08:31:58.820607Z","iopub.status.idle":"2025-01-18T08:31:59.964331Z","shell.execute_reply.started":"2025-01-18T08:31:58.820568Z","shell.execute_reply":"2025-01-18T08:31:59.963358Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class CNNBlock(nn.Module):\n    def __init__(self, in_channels, out_channels, kernel_size=4, stride=2, device=\"cuda\", use_dropout=True, norm_method = 'batch'):\n        \"\"\"\n        Initialize the CNNBlock with convolution, batch normalization, and activation layers.\n        \n        Args:\n            in_channels (int): Number of input channels.\n            out_channels (int): Number of output channels.\n            stride (int, optional): Stride for the convolution. Defaults to 2.\n            device (str, optional): Device to run the operations on ('cuda' or 'cpu'). Defaults to \"cuda\".\n        \"\"\"\n        super().__init__()  # Initialize the parent class\n        self.device = device  # Set the device\n        layers = [\n            nn.Conv2d(\n                in_channels, \n                out_channels, \n                kernel_size=kernel_size, \n                stride=stride, \n                padding=1,\n                padding_mode='reflect', \n                bias=False,\n                device=self.device\n            ),  # Convolutional layer with specified parameters\n            nn.BatchNorm2d(out_channels, device=device) if norm_method == 'batch' else \n            nn.InstanceNorm2d(out_channels, device=device) if norm_method == 'instance' else \n            None\n            ,  \n            nn.LeakyReLU(0.2),  # Leaky ReLU activation with negative slope of 0.2\n            ]\n        self.seq = nn.Sequential(*[layer for layer in layers if layer is not None])\n        self.use_dropout = use_dropout  # Set whether to use dropout\n        self.dropout = nn.Dropout(0.5)  # Define a dropout layer with 50% dropout rate\n    \n        \n    def forward(self, x):\n        \"\"\"\n        Forward pass through the CNNBlock.\n        \n        Args:\n            x (torch.Tensor): Input tensor.\n        \n        Returns:\n            torch.Tensor: Output tensor after convolution, batch normalization, and activation.\n        \"\"\"\n        x = self.seq(x)  # Apply convolution, normalization, and activation\n        return self.dropout(x) if self.use_dropout else x  # Apply dropout if enabled, else return the tensor\n\n\nclass Discriminator(nn.Module):\n    def __init__(self, in_channels=3, features=[64, 128, 256], device=\"cuda\"):\n        \"\"\"\n        Initialize the Discriminator model using a series of CNNBlocks.\n        \n        Args:\n            in_channels (int, optional): Number of input channels. Defaults to 3 (e.g., RGB images).\n            features (list, optional): List of feature sizes for each CNNBlock. Defaults to [64, 128, 256, 512].\n            device (str, optional): Device to run the operations on ('cuda' or 'cpu'). Defaults to \"cuda\".\n        \"\"\"\n        super().__init__()  # Initialize the parent class\n        self.device = device  # Set the device\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                device=self.device\n            ),  # Initial convolutional layer\n            nn.LeakyReLU(0.2)  # Leaky ReLU activation\n        )\n        \n        layers = []  # Initialize a list to hold subsequent CNNBlocks\n        in_channels = features[0]  # Set the initial number of input channels\n        for feature in features[1:]:\n            layers.append(\n                CNNBlock(in_channels, feature, kernel_size=4, stride=1 if feature == features[-1] else 2, device=self.device, use_dropout=True),\n            )  # Append CNNBlock with appropriate stride\n            in_channels = feature  # Update the number of input channels for the next block\n        \n        layers.append(\n            nn.Conv2d(\n                in_channels, \n                1, \n                kernel_size=4, \n                stride=2, \n                padding=1, \n                padding_mode='reflect', \n                device=self.device\n            )\n        )  # Final convolutional layer to produce a single output channel (real/fake classification)\n        \n        self.model = nn.Sequential(*layers)  # Combine all layers into a sequential model\n\n        def weights_init_normal(m):\n            if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d) or isinstance(m, nn.Linear):\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.0)\n                    \n        self.apply(weights_init_normal)\n        \n    def forward(self, x):\n        \"\"\"\n        Forward pass through the Discriminator.\n        \n        Args:\n            x (torch.Tensor): Input tensor (e.g., an image).\n        \n        Returns:\n            torch.Tensor: Output tensor representing the discriminator's prediction.\n        \"\"\"\n        x = self.initial(x)  # Apply the initial convolution and activation\n        return self.model(x)  # Apply the subsequent CNNBlocks and final convolution\n        \n\nclass Interpolate(nn.Module):\n    def __init__(self, scale_factor=None, mode='linear'):\n        \"\"\"\n        Initializes the Interpolate module.\n\n        Args:\n            size (int or tuple, optional): The output size. Can be a single integer or a tuple of integers.\n            scale_factor (float or tuple, optional): The multiplier for spatial size. Can be a single float or a tuple of floats.\n            mode (str, optional): The upsampling mode to use. Options include 'nearest', 'linear', 'bilinear', 'bicubic', and 'trilinear'.\n            align_corners (bool, optional): Geometrically, if True, the corner pixels of the input and output tensors are aligned.\n        \"\"\"\n        super(Interpolate, self).__init__()\n        self.scale_factor = scale_factor\n        self.mode = mode\n\n    def forward(self, x):\n        \"\"\"\n        Forward pass for the Interpolate module.\n\n        Args:\n            x (torch.Tensor): Input tensor.\n\n        Returns:\n            torch.Tensor: Interpolated output tensor.\n        \"\"\"\n        if x.dim() == 4:  # 4D tensor\n            n, c, h, w = x.shape\n        elif x.dim() == 3:  # 3D tensor\n            c, h, w = x.shape\n            n = 1  # Default batch size for 3D tensor\n        else:\n            raise ValueError(f\"Unexpected tensor shape: {x.shape}\")\n            \n        return F.interpolate(x.view(n, c, h, w), scale_factor=self.scale_factor, mode=self.mode)\n\n\nclass ResizeConv(nn.Module):\n    def __init__(self, in_channels, out_channels, scale_factor, kernel_size, stride, padding, device):\n        super(ResizeConv, self).__init__()\n        self.scale_factor = scale_factor\n        self.conv = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride, padding=padding, device=device, bias=False)\n        self.interpolate = Interpolate(scale_factor=[self.scale_factor, self.scale_factor], mode='bilinear')\n\n    def forward(self, x):\n        x = self.conv(self.interpolate(x))\n        return x\n        \n        \nclass Block(nn.Module):\n    def __init__(self, in_channels, out_channels, down=True, act='relu', use_dropout=False, device=\"cuda\"):\n        \"\"\"\n        Initialize a Block, which can act as either a downsampling or upsampling layer.\n        \n        Args:\n            in_channels (int): Number of input channels.\n            out_channels (int): Number of output channels.\n            down (bool, optional): If True, performs downsampling using Conv2d; else, upsampling using ConvTranspose2d. Defaults to True.\n            act (str, optional): Activation function to use ('relu' or 'leaky'). Defaults to 'relu'.\n            use_dropout (bool, optional): Whether to include a dropout layer. Defaults to False.\n            device (str, optional): Device to run the operations on ('cuda' or 'cpu'). Defaults to \"cuda\".\n        \"\"\"\n        super().__init__()  # Initialize the parent class\n        self.device = device  # Set the device\n        self.seq = nn.Sequential(\n            nn.Conv2d(\n                in_channels, \n                out_channels, \n                kernel_size=4, \n                stride=2, \n                padding=1,\n                bias=False,\n                padding_mode='reflect', \n                device=self.device\n            ) if down else ResizeConv(\n                in_channels, \n                out_channels,\n                scale_factor=2.0,\n                kernel_size=3, \n                stride=1, \n                padding=1,\n                device=self.device\n            ),  # Choose Conv2d for downsampling or ConvTranspose2d for upsampling\n            nn.InstanceNorm2d(out_channels, device=self.device),  # Batch normalization\n            nn.ReLU() if act == 'relu' else nn.LeakyReLU(0.2),  # Activation function\n        )\n        self.use_dropout = use_dropout  # Set whether to use dropout\n        self.dropout = nn.Dropout(0.5)  # Define a dropout layer with 50% dropout rate\n    \n    def forward(self, x):\n        \"\"\"\n        Forward pass through the Block.\n        \n        Args:\n            x (torch.Tensor): Input tensor.\n        \n        Returns:\n            torch.Tensor: Output tensor after convolution, normalization, activation, and optional dropout.\n        \"\"\"\n        x = self.seq(x)  # Apply convolution, normalization, and activation\n        return self.dropout(x) if self.use_dropout else x  # Apply dropout if enabled, else return the tensor\n\n\nclass Generator(nn.Module):\n    def __init__(self, in_channels=3, features=64, device=\"cuda\"):\n        \"\"\"\n        Initialize the Generator model using a series of downsampling and upsampling Blocks.\n        \n        Args:\n            in_channels (int, optional): Number of input channels. Defaults to 3 (e.g., RGB images).\n            features (int, optional): Base number of feature channels. Defaults to 64.\n            device (str, optional): Device to run the operations on ('cuda' or 'cpu'). Defaults to \"cuda\".\n        \"\"\"\n        super().__init__()  # Initialize the parent class\n        self.device = device  # Set the device\n        self.initial_down = nn.Sequential(\n            nn.Conv2d(\n                in_channels, \n                features, \n                kernel_size=4, \n                stride=2, \n                padding=1, \n                padding_mode='reflect', \n                device=self.device\n            ),  # Initial convolutional layer\n            nn.LeakyReLU(0.2),  # Leaky ReLU activation\n        )  # Output size: 128\n        \n        # Define downsampling Blocks\n        self.down1 = Block(features, features*2, down=True, act='leaky', use_dropout=False)   # Output size: 64\n        self.down2 = Block(features*2, features*4, down=True, act='leaky', use_dropout=False) # Output size: 32\n        self.down3 = Block(features*4, features*8, down=True, act='leaky', use_dropout=False) # Output size: 16\n        self.down4 = Block(features*8, features*8, down=True, act='leaky', use_dropout=False) # Output size: 8\n        self.down5 = Block(features*8, features*8, down=True, act='leaky', use_dropout=False) # Output size: 4\n        self.down6 = Block(features*8, features*8, down=True, act='leaky', use_dropout=False) # Output size: 2\n        \n        # Bottleneck layer\n        self.bottleneck = nn.Sequential(\n            nn.Conv2d(\n                features*8, \n                features*8, \n                kernel_size=4, \n                stride=2, \n                padding=1, \n                padding_mode='reflect', \n                device=self.device\n            ),  # Convolutional layer at the bottleneck\n            nn.ReLU(),  # ReLU activation\n        )  # Output size: 1x1\n        \n        # Define upsampling Blocks with skip connections\n        self.up1 = Block(features*8, features*8, down=False, act='relu', use_dropout=True)  # Output size: 2\n        self.up2 = Block(features*8*2, features*8, down=False, act='relu', use_dropout=True)  # Output size: 4\n        self.up3 = Block(features*8*2, features*8, down=False, act='relu', use_dropout=True)  # Output size: 8\n        self.up4 = Block(features*8*2, features*8, down=False, act='relu', use_dropout=False) # Output size: 16\n        self.up5 = Block(features*8*2, features*4, down=False, act='relu', use_dropout=False) # Output size: 32\n        self.up6 = Block(features*4*2, features*2, down=False, act='relu', use_dropout=False) # Output size: 64\n        self.up7 = Block(features*2*2, features, down=False, act='relu', use_dropout=False)   # Output size: 128\n        \n        # Final upsampling layer to reconstruct the image\n        self.finil_up = nn.Sequential(\n            ResizeConv(\n                features*2, \n                features, \n                scale_factor=2.0,\n                kernel_size=3,\n                stride=1, \n                padding=1,\n                device=self.device\n            ),\n            CNNBlock(features, in_channels, kernel_size=3, stride=1, device=device, norm_method = 'none'),\n            nn.Tanh(),  # Tanh activation to scale the output between -1 and 1\n        )\n        \n        def weights_init_normal(m):\n            if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d) or isinstance(m, nn.Linear):\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.0)\n                    \n        self.apply(weights_init_normal)\n        \n    def forward(self, x):\n        \"\"\"\n        Forward pass through the Generator, implementing the U-Net architecture with skip connections.\n        \n        Args:\n            x (torch.Tensor): Input tensor (e.g., an image).\n        \n        Returns:\n            torch.Tensor: Output tensor representing the generated image.\n        \"\"\"\n        d1 = self.initial_down(x)  # Apply initial downsampling\n        d2 = self.down1(d1)  # Apply first downsampling Block\n        d3 = self.down2(d2)  # Apply second downsampling Block\n        d4 = self.down3(d3)  # Apply third downsampling Block\n        d5 = self.down4(d4)  # Apply fourth downsampling Block\n        d6 = self.down5(d5)  # Apply fifth downsampling Block\n        d7 = self.down6(d6)  # Apply sixth downsampling Block\n        \n        bottleneck = self.bottleneck(d7)  # Apply bottleneck layer\n        \n        up1 = self.up1(bottleneck)  # Apply first upsampling Block\n        up2 = self.up2(torch.cat([up1, d7], 1))  # Concatenate with corresponding downsampled feature map and apply second upsampling Block\n        up3 = self.up3(torch.cat([up2, d6], 1))  # Concatenate with corresponding downsampled feature map and apply third upsampling Block\n        up4 = self.up4(torch.cat([up3, d5], 1))  # Concatenate with corresponding downsampled feature map and apply fourth upsampling Block\n        up5 = self.up5(torch.cat([up4, d4], 1))  # Concatenate with corresponding downsampled feature map and apply fifth upsampling Block\n        up6 = self.up6(torch.cat([up5, d3], 1))  # Concatenate with corresponding downsampled feature map and apply sixth upsampling Block\n        up7 = self.up7(torch.cat([up6, d2], 1))  # Concatenate with corresponding downsampled feature map and apply seventh upsampling Block\n\n        # Concatenate with initial downsampled feature map and apply final upsampling to generate the output image\n        return self.finil_up(torch.cat([up7, d1], 1))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-18T08:31:59.965804Z","iopub.execute_input":"2025-01-18T08:31:59.966086Z","iopub.status.idle":"2025-01-18T08:31:59.997733Z","shell.execute_reply.started":"2025-01-18T08:31:59.966057Z","shell.execute_reply":"2025-01-18T08:31:59.996934Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class Trainer:\n    def __init__(\n        self,\n        train_data: DataLoader,\n        val_data: DataLoader,\n        target_generator: torch.nn.Module,\n        target_discriminator: torch.nn.Module,\n        source_generator: torch.nn.Module,\n        source_discriminator: torch.nn.Module,\n        nb_epochs: int = 5,\n        device: str = \"cuda\",\n        save_path: str = None,\n        discriminator_workouts_per_step: int = 1,\n        generator_workouts_per_step: int = 1,\n        lambda_coef: int = 10,\n        buffer_size: int = 64,\n        load_m: bool = True,\n    ):\n        \"\"\"\n        Initialize the Trainer with data loaders, models, training parameters, and device configuration.\n        \n        Args:\n            train_data (DataLoader): DataLoader for training data.\n            val_data (DataLoader): DataLoader for validation data.\n            generator (torch.nn.Module): Generator model.\n            discriminator (torch.nn.Module): Discriminator model.\n            nb_epochs (int, optional): Number of training epochs. Defaults to 5.\n            device (str, optional): Device to run the training on ('cuda' or 'cpu'). Defaults to \"cuda\".\n            save_path (str, optional): Path to save the trained models. Defaults to None.\n        \"\"\"\n        self.train_data = train_data                      # Assign training data loader\n        self.val_data = val_data                          # Assign validation data loader\n        self.target_generator = target_generator          # Assign generator model\n        self.target_discriminator = target_discriminator  # Assign discriminator model\n        self.source_generator = source_generator          # Assign generator model\n        self.source_discriminator = source_discriminator  # Assign discriminator model\n        self.nb_epochs = nb_epochs                        # Set number of epochs\n        self.device = device                              # Set device for computation\n        self.save_path = Path(save_path) if save_path else save_path  # Set save path as Path object if provided\n\n        self.discriminator_workouts_per_step = discriminator_workouts_per_step\n        self.generator_workouts_per_step = generator_workouts_per_step\n\n        self.lambda_coef = lambda_coef\n\n        self.EXPONENTIAL_LOSS = 0.1\n\n        self.buffer_size = buffer_size\n        self.t_buffer_images = None\n        self.s_buffer_images = None\n\n        self.load_m = load_m\n        self.scaler = GradScaler()\n\n        # Initialize a fixed set of validation samples for monitoring generator progress\n        self.z = next(iter(self.val_data))[0][:32].to(self.device)\n        \n        # Initialize a dictionary to store training logs\n        self.logs = {\n            \"Step\": [],\n            \"Train_tg_loss\": [],\n            \"Train_td_loss\": [],\n            \"Val_tg_loss\": [],\n            \"Val_td_loss\": [],\n            \"Train_sg_loss\": [],\n            \"Train_sd_loss\": [],\n            \"Val_sg_loss\": [],\n            \"Val_sd_loss\": [],\n            \"Samples\": [],\n        }\n\n    def init_optimizers(self, lr: float=1e-5, betas: tuple=(0.5, 0.999), weight_decay: float=1e-5):\n        \"\"\"\n        Initialize the optimizers for the generator and discriminator.\n        \n        Args:\n            lr (float, optional): Learning rate for the optimizers. Defaults to 1e-5.\n            betas (tuple, optional): Beta parameters for the AdamW optimizer. Defaults to (0.5, 0.999).\n            weight_decay (float, optional): Weight Decay regularization. Defaults to 1e-5\n        \"\"\"\n        # Initialize Adam optimizer for the generator\n        self.tg_optimizer = torch.optim.AdamW(\n            self.target_generator.parameters(), lr=lr, betas=betas, weight_decay=weight_decay\n        )\n        # Initialize Adam optimizer for the discriminator\n        self.td_optimizer = torch.optim.AdamW(\n            self.target_discriminator.parameters(), lr=lr, betas=betas, weight_decay=weight_decay\n        )\n        # Initialize Adam optimizer for the generator\n        self.sg_optimizer = torch.optim.AdamW(\n            self.source_generator.parameters(), lr=lr, betas=betas, weight_decay=weight_decay\n        )\n        # Initialize Adam optimizer for the discriminator\n        self.sd_optimizer = torch.optim.AdamW(\n            self.source_discriminator.parameters(), lr=lr, betas=betas, weight_decay=weight_decay\n        )\n\n    def init_buffers(self):\n        with torch.no_grad():\n            initial_images = []\n            initial_mones = []\n            for batch in self.train_data:\n                X, y = batch\n                initial_images.append(X)\n                initial_mones.append(y)\n                if len(initial_images) >= self.buffer_size:\n                    break\n                    \n            self.t_buffer_images = self.target_generator(torch.cat(initial_images, dim=0).to(self.device)[:self.buffer_size]).to(torch.float16)\n            self.s_buffer_images = self.source_generator(torch.cat(initial_mones, dim=0).to(self.device)[:self.buffer_size]).to(torch.float16)\n\n    def init_epoch_logs(self):\n        self.epoch_logs = {\n            \"Train_tg_loss\": [],\n            \"Train_td_loss\": [],\n            \"Val_tg_loss\": [],\n            \"Val_td_loss\": [],\n            \"Train_sg_loss\": [],\n            \"Train_sd_loss\": [],\n            \"Val_sg_loss\": [],\n            \"Val_sd_loss\": [],\n        }\n    \n    def update_buffers(self, x, y, size):\n        indices = torch.randint(high = self.buffer_size, size = (int(np.random.binomial(size, 0.5)), ))\n        self.t_buffer_images[indices] = x[indices].detach()\n        self.s_buffer_images[indices] = y[indices].detach()\n\n    def sample_buffers(self, size):\n        indices = torch.randint(high = self.buffer_size, size = size)\n        return self.t_buffer_images[indices], self.s_buffer_images[indices]\n\n    def scaler_grad(self, optim, params):\n        self.scaler.unscale_(optim)\n        torch.nn.utils.clip_grad_norm_(params, max_norm=2.0)\n        self.scaler.step(optim)\n    \n    def plot_losses(self):\n        (train_tg_loss, train_td_loss, train_sg_loss, train_sd_loss,\n         val_tg_loss, val_td_loss, val_sg_loss, val_sd_loss) = (\n            self.epoch_logs[\"Train_tg_loss\"],\n            self.epoch_logs[\"Train_td_loss\"],\n            self.epoch_logs[\"Train_sg_loss\"],\n            self.epoch_logs[\"Train_sd_loss\"],\n            self.epoch_logs[\"Val_tg_loss\"],\n            self.epoch_logs[\"Val_td_loss\"],\n            self.epoch_logs[\"Val_sg_loss\"],\n            self.epoch_logs[\"Val_sd_loss\"]\n        )\n        \n        t_steps_per_epochs = range(1, len(train_tg_loss) + 1)\n        v_steps_per_epochs = range(1, len(val_tg_loss) + 1)\n        # Create subplots for training and validation losses side by side\n        fig, axes = plt.subplots(1, 2, figsize=(16, 6), sharey=True)\n\n        axes[0].plot(t_steps_per_epochs, train_tg_loss, label=\"Train TG Loss\", marker='o')\n        axes[0].plot(t_steps_per_epochs, train_td_loss, label=\"Train TD Loss\", marker='s')\n        axes[0].plot(t_steps_per_epochs, train_sg_loss, label=\"Train SG Loss\", marker='^')\n        axes[0].plot(t_steps_per_epochs, train_sd_loss, label=\"Train SD Loss\", marker='x')\n        axes[0].set_title(\"Training Losses Over Steps Per Epoch\", fontsize=16)\n        axes[0].set_xlabel(\"Steps Per Epoch\", fontsize=14)\n        axes[0].set_ylabel(\"Loss\", fontsize=14)\n        axes[0].set_ylim(0, 2.5)\n        axes[0].legend(fontsize=12)\n        axes[0].grid(True, linestyle='--', alpha=0.6)\n\n        axes[1].plot(v_steps_per_epochs, val_tg_loss, label=\"Val TG Loss\", marker='o')\n        axes[1].plot(v_steps_per_epochs, val_td_loss, label=\"Val TD Loss\", marker='s')\n        axes[1].plot(v_steps_per_epochs, val_sg_loss, label=\"Val SG Loss\", marker='^')\n        axes[1].plot(v_steps_per_epochs, val_sd_loss, label=\"Val SD Loss\", marker='x')\n        axes[1].set_title(\"Validation Losses Over Steps Per Epoch\", fontsize=16)\n        axes[1].set_xlabel(\"Steps Per Epoch\", fontsize=14)\n        axes[1].set_ylim(0, 2.5)\n        axes[1].legend(fontsize=12)\n        axes[1].grid(True, linestyle='--', alpha=0.6)\n\n        # Adjust layout and show the plots\n        plt.tight_layout()\n        plt.show()\n\n    def plot_images(self):\n        fake_img = self.target_generator((photos[0].unsqueeze(0)).to('cuda'))\n        cycle_img = self.source_generator(fake_img)\n                \n        fake_img = fake_img[0].cpu().detach()\n        cycle_img = cycle_img[0].cpu().detach()\n                \n        # Plot the original photo and the generated fake image\n        fig, ax = plt.subplots(1, 3, figsize=(20, 4))\n        ax[0].imshow(photos[0].permute(1, 2, 0).numpy() * 0.5 + 0.5)\n        ax[0].set_title(\"Photo\")\n        ax[0].set(xticks=[], yticks=[])\n        ax[1].imshow(fake_img.permute(1, 2, 0).numpy() * 0.5 + 0.5)\n        ax[1].set_title(\"FakeMonet\")\n        ax[1].set(xticks=[], yticks=[])\n        ax[2].imshow(cycle_img.permute(1, 2, 0).numpy() * 0.5 + 0.5)\n        ax[2].set_title(\"CyclePhoto\")\n        ax[2].set(xticks=[], yticks=[])\n        plt.show()\n\n    \n    def train(self):\n        \"\"\"\n        Execute the training loop for the GAN, including training and validation phases,\n        logging, and model checkpointing.\n        \n        Returns:\n            dict: Dictionary containing training logs.\n        \"\"\"\n        # Ensure that both optimizers are initialized before training\n        assert (self.tg_optimizer is not None) and (\n            self.td_optimizer is not None\n        ) and (self.sg_optimizer is not None) and (\n            self.sd_optimizer is not None\n        ), \"Please run Trainer().init_optimizer()\"\n\n        if self.load_m:\n            # If a saved model exists at the save path, load the model weights\n            if self.save_path and self.save_path.exists():\n                self.load_model()\n            \n        best_score = torch.inf  # Initialize the best validation score to infinity\n\n        # Iterate over each epoch\n        for i in range(self.nb_epochs):\n            self.init_epoch_logs()\n            # Initialize cumulative losses for the epoch\n            train_td_loss, train_tg_loss, val_td_loss, val_tg_loss = 0, 0, 0, 0\n            train_sd_loss, train_sg_loss, val_sd_loss, val_sg_loss = 0, 0, 0, 0\n            train_cycle_loss, val_cycle_loss = 0, 0\n            train_identity_loss, val_identity_loss = 0, 0\n            \n            self.target_generator.train()      # Set generator to training mode\n            self.target_discriminator.train()  # Set discriminator to training mode\n            self.source_generator.train()      # Set generator to training mode\n            self.source_discriminator.train()  # Set discriminator to training mode\n\n            # Training loop with progress bar\n            loop = tqdm(\n                enumerate(itertools.islice(self.train_data, len(self.train_data) - 1)),  # Skip the last element\n                desc=f\"Epoch {i + 1}/{self.nb_epochs} training\",\n                leave=True,\n                total=len(self.train_data) - 1,  # Update the total to reflect the adjusted length\n            )\n            for step, (x, y) in loop:\n                x = x.to(self.device)  # Move input data to the specified device\n                y = y.to(self.device)  # Move target data to the specified device\n            \n                # Perform a training step and retrieve generator and discriminator losses\n                with autocast():\n                    tg_loss, td_loss, sg_loss, sd_loss, cycle_loss, identity_loss = self.train_step(x, y)\n            \n                # Accumulate the losses with exponential averaging\n                train_tg_loss = (1 - self.EXPONENTIAL_LOSS) * train_tg_loss + self.EXPONENTIAL_LOSS * tg_loss\n                train_td_loss = (1 - self.EXPONENTIAL_LOSS) * train_td_loss + self.EXPONENTIAL_LOSS * td_loss\n                train_sg_loss = (1 - self.EXPONENTIAL_LOSS) * train_sg_loss + self.EXPONENTIAL_LOSS * sg_loss\n                train_sd_loss = (1 - self.EXPONENTIAL_LOSS) * train_sd_loss + self.EXPONENTIAL_LOSS * sd_loss\n                train_cycle_loss = (1 - self.EXPONENTIAL_LOSS) * train_cycle_loss + self.EXPONENTIAL_LOSS * cycle_loss\n                train_identity_loss = (1 - self.EXPONENTIAL_LOSS) * train_identity_loss + self.EXPONENTIAL_LOSS * identity_loss\n            \n                # Update the progress bar with average losses\n                loop.set_postfix_str(\n                    f\"tg_loss: {train_tg_loss :.2f}, \"\n                    f\"td_loss: {train_td_loss :.2f} :: \"\n                    f\"sg_loss: {train_sg_loss :.2f}, \"\n                    f\"sd_loss: {train_sd_loss :.2f} :: \"\n                    f\"cycle_loss: {train_cycle_loss :.2f}, \"\n                    f\"identity_loss: {train_identity_loss :.2f}\"\n                )\n            \n            # Validation phase\n            self.target_generator.eval()      # Set generator to evaluation mode\n            self.target_discriminator.eval()  # Set discriminator to evaluation mode\n            self.source_generator.eval()      # Set generator to evaluation mode\n            self.source_discriminator.eval()  # Set discriminator to evaluation mode\n            \n            # Validation loop with progress bar\n            loop = tqdm(\n                enumerate(self.val_data),\n                desc=f\"Epoch {i + 1}/{self.nb_epochs} validation\",\n                leave=True,\n                total=len(self.val_data),\n            )\n            for step, (x, y) in loop:\n                x = x.to(self.device)  # Move input data to the specified device\n                y = y.to(self.device)  # Move target data to the specified device\n            \n                # Perform a validation step and retrieve generator and discriminator losses\n                with autocast():\n                    tg_loss, td_loss, sg_loss, sd_loss, cycle_loss, identity_loss = self.val_step(x, y)\n            \n                # Accumulate the validation losses with exponential averaging\n                val_tg_loss = (1 - self.EXPONENTIAL_LOSS) * val_tg_loss + self.EXPONENTIAL_LOSS * tg_loss\n                val_td_loss = (1 - self.EXPONENTIAL_LOSS) * val_td_loss + self.EXPONENTIAL_LOSS * td_loss\n                val_sg_loss = (1 - self.EXPONENTIAL_LOSS) * val_sg_loss + self.EXPONENTIAL_LOSS * sg_loss\n                val_sd_loss = (1 - self.EXPONENTIAL_LOSS) * val_sd_loss + self.EXPONENTIAL_LOSS * sd_loss\n                val_cycle_loss = (1 - self.EXPONENTIAL_LOSS) * val_cycle_loss + self.EXPONENTIAL_LOSS * cycle_loss\n                val_identity_loss = (1 - self.EXPONENTIAL_LOSS) * val_identity_loss + self.EXPONENTIAL_LOSS * identity_loss\n            \n                # Update the progress bar with average losses\n                loop.set_postfix_str(\n                    f\"tg_loss: {val_tg_loss :.2f}, \"\n                    f\"td_loss: {val_td_loss :.2f} :: \"\n                    f\"sg_loss: {val_sg_loss :.2f}, \"\n                    f\"sd_loss: {val_sd_loss :.2f} :: \"\n                    f\"cycle_loss: {val_cycle_loss :.2f}, \"\n                    f\"identity_loss: {val_identity_loss :.2f}\"\n                )\n\n            # Check if the current validation loss is the best so far\n            if self.save_path and best_score > val_tg_loss:\n                best_score = val_tg_loss  # Update the best score\n                self.save_model()  # Save the current model as the best model\n\n            # Log the metrics for the current epoch\n            self.log_metrics(\n                step=i,\n                train_tg_loss=train_tg_loss,\n                train_td_loss=train_td_loss,\n                val_tg_loss=val_tg_loss,\n                val_td_loss=val_td_loss,\n                train_sg_loss=train_sg_loss,\n                train_sd_loss=train_sd_loss,\n                val_sg_loss=val_sg_loss,\n                val_sd_loss=val_sd_loss,\n            )\n            # Generate and display sample images from the generator\n            self.plot_images()\n            self.plot_losses()\n            \n            \n        return self.logs  # Return the training logs\n\n    def train_step(\n        self,\n        x: torch.Tensor,\n        y: torch.Tensor\n    ) -> tuple:\n        \"\"\"\n        Perform a single training step for both generator and discriminator.\n        \n        Args:\n            x (torch.Tensor): Input tensor (e.g., photos).\n            y (torch.Tensor): Target tensor (e.g., Monet paintings).\n            \n        Returns:\n            tuple: Generator loss and discriminator loss.\n        \"\"\"               \n        td_loss, tg_loss, sd_loss, sg_loss = -1, -1, -1, -1\n        for _ in range(self.discriminator_workouts_per_step):\n            self.td_optimizer.zero_grad(set_to_none=True)  # Reset discriminator gradients\n            self.sd_optimizer.zero_grad(set_to_none=True)  # Reset discriminator gradients\n\n            # Generate fake images and classify them with the discriminator\n            # Classify real images with the discriminator\n\n            # Merge and update sampler\n\n            tfake, sfake = self.target_generator(x), self.source_generator(y)\n            fakey, fakex = self.sample_buffers(size = (BATCH_SIZE, ))\n            self.update_buffers(x = sfake, y = tfake, size = (BATCH_SIZE, ))\n        \n            tfake = torch.cat([tfake, fakey], dim=0)\n            sfake = torch.cat([sfake, fakex], dim=0)\n            \n            tfake, treal  = self.target_discriminator(tfake), self.target_discriminator(y)\n            sfake, sreal  = self.source_discriminator(sfake), self.source_discriminator(x)\n            \n            # Calculate discriminator loss: real images should be classified as ones and fake as zeros\n            td_loss = (\n                F.mse_loss(treal, torch.ones_like(treal, device=self.device)) + \n                F.mse_loss(tfake, torch.zeros_like(tfake, device=self.device))\n            )\n            sd_loss = (\n                F.mse_loss(sreal, torch.ones_like(sreal, device=self.device)) + \n                F.mse_loss(sfake, torch.zeros_like(sfake, device=self.device))\n            )\n\n            self.scaler.scale(td_loss).backward()\n            self.scaler.scale(sd_loss).backward()\n            \n            self.scaler_grad(self.td_optimizer, self.target_discriminator.parameters())\n            self.scaler_grad(self.sd_optimizer, self.source_discriminator.parameters())\n\n            self.scaler.update()\n\n        # Re-generate fake images for generator training\n        for _ in range(self.generator_workouts_per_step):\n            self.tg_optimizer.zero_grad(set_to_none=True)  # Reset generator gradients\n            self.sg_optimizer.zero_grad(set_to_none=True)  # Reset generator gradients\n\n            # Identity Generation\n            iden_x = self.target_generator(x)\n            iden_y = self.source_generator(y)\n            \n            tfake, sfake = self.target_generator(x), self.source_generator(y)\n            tcycle, scycle = self.target_generator(sfake), self.source_generator(tfake)\n            tfake, sfake = self.target_discriminator(tfake), self.source_discriminator(sfake)\n            \n            # Calculate generator loss: including adversary and cycle loss\n            tg_loss = F.mse_loss(tfake, torch.ones_like(tfake, device=self.device))\n            sg_loss = F.mse_loss(sfake, torch.ones_like(sfake, device=self.device))\n            cycle_loss = F.l1_loss(scycle, x) + F.l1_loss(tcycle, y)\n            identity_loss = F.mse_loss(iden_x ,x) + F.mse_loss(iden_y, y)\n            \n            generators_loss = (\n                tg_loss + sg_loss +\n                self.lambda_coef * cycle_loss +\n                self.lambda_coef * 0.5 * identity_loss\n            )\n            \n            self.scaler.scale(generators_loss).backward()  # Backpropagate generator loss\n            self.scaler_grad(self.tg_optimizer, self.target_generator.parameters())\n            self.scaler_grad(self.sg_optimizer, self.source_generator.parameters())\n            self.scaler.update()\n\n\n        self.epoch_log_metrics(\n            tg_loss=tg_loss.item(),\n            td_loss=td_loss.item(),\n            sg_loss=sg_loss.item(),\n            sd_loss=sd_loss.item(),\n            train=True,\n        )\n\n        return tg_loss.item(), td_loss.item(), sg_loss.item(), sd_loss.item(), cycle_loss.item(), identity_loss.item()  # Return scalar losses\n\n    @torch.no_grad()\n    def val_step(\n        self,\n        x: torch.Tensor,\n        y: torch.Tensor,\n    ) -> tuple:\n        \"\"\"\n        Perform a single validation step without updating the model.\n        \n        Args:\n            x (torch.Tensor): Input tensor (e.g., photos).\n            y (torch.Tensor): Target tensor (e.g., Monet paintings).\n            \n        Returns:\n            tuple: Generator loss and discriminator loss.\n        \"\"\"\n        # Identity Generation\n        iden_x = self.target_generator(x)\n        iden_y = self.source_generator(y)\n        \n        # Generate fake images and classify them with the discriminator\n        tfake = self.target_generator(x)\n        sfake = self.source_generator(y)\n        tcycle = self.target_generator(sfake)\n        scycle = self.source_generator(tfake)\n        \n        # Classify real images with the discriminator\n        treal = self.target_discriminator(y)\n        sreal = self.source_discriminator(x)\n        tfake = self.target_discriminator(tfake)\n        sfake = self.source_discriminator(sfake)\n        \n        # Calculate generator loss: fake images should be classified as ones\n        tg_loss = torch.nn.functional.mse_loss(tfake, torch.ones_like(tfake, device=self.device))\n        sg_loss = torch.nn.functional.mse_loss(sfake, torch.ones_like(sfake, device=self.device))\n        # Calculate discriminator loss: real images as ones and fake as zeros\n        td_loss = (\n            torch.nn.functional.mse_loss(treal, torch.ones_like(treal, device=self.device)) + \n            torch.nn.functional.mse_loss(tfake, torch.zeros_like(tfake, device=self.device))\n        )\n        sd_loss = (\n            torch.nn.functional.mse_loss(sreal, torch.ones_like(sreal, device=self.device)) + \n            torch.nn.functional.mse_loss(sfake, torch.zeros_like(sfake, device=self.device))\n        )\n        cycle_loss = F.l1_loss(scycle, x) + F.l1_loss(tcycle, y)\n        identity_loss = F.mse_loss(iden_x, x) + F.mse_loss(iden_y, y)\n\n        self.epoch_log_metrics(\n            tg_loss=tg_loss.item(),\n            td_loss=td_loss.item(),\n            sg_loss=sg_loss.item(),\n            sd_loss=sd_loss.item(),\n            train=False,\n        )\n        \n        return tg_loss.item(), td_loss.item(), sg_loss.item(), sd_loss.item(), cycle_loss.item(), identity_loss.item()  # Return scalar losses\n\n    @torch.no_grad() # Make this one function \n    def epoch_log_metrics(\n        self,\n        tg_loss: torch.Tensor,\n        td_loss: torch.Tensor,\n        sg_loss: torch.Tensor,\n        sd_loss: torch.Tensor,\n        train: bool,\n    ):\n        if train:\n            self.epoch_logs[\"Train_tg_loss\"].append(tg_loss)      # Log training generator loss\n            self.epoch_logs[\"Train_td_loss\"].append(td_loss)      # Log training discriminator loss\n            self.epoch_logs[\"Train_sg_loss\"].append(sg_loss)      # Log training generator loss\n            self.epoch_logs[\"Train_sd_loss\"].append(sd_loss)      # Log training discriminator loss\n        else:\n            self.epoch_logs[\"Val_tg_loss\"].append(tg_loss)        # Log validation generator loss\n            self.epoch_logs[\"Val_sg_loss\"].append(sg_loss)        # Log validation generator loss\n            self.epoch_logs[\"Val_sd_loss\"].append(sd_loss)        # Log validation discriminator loss\n            self.epoch_logs[\"Val_td_loss\"].append(td_loss)        # Log validation discriminator loss\n\n    @torch.no_grad()\n    def log_metrics(\n        self,\n        step: int,\n        train_tg_loss: torch.Tensor,\n        train_td_loss: torch.Tensor,\n        val_tg_loss: torch.Tensor,\n        val_td_loss: torch.Tensor,\n        train_sg_loss: torch.Tensor,\n        train_sd_loss: torch.Tensor,\n        val_sg_loss: torch.Tensor,\n        val_sd_loss: torch.Tensor,\n    ):\n        \"\"\"\n        Log the training and validation metrics for the current epoch.\n        \n        Args:\n            step (int): Current epoch number.\n            train_g_loss (torch.Tensor): EMA generator loss during training.\n            train_d_loss (torch.Tensor): Cumulative discriminator loss during training.\n            val_g_loss (torch.Tensor): Cumulative generator loss during validation.\n            val_d_loss (torch.Tensor): Cumulative discriminator loss during validation.\n        \"\"\"\n        self.logs[\"Step\"].append(step)  # Log the current epoch\n        self.logs[\"Train_tg_loss\"].append(train_tg_loss)    # Log average training generator loss\n        self.logs[\"Train_td_loss\"].append(train_td_loss)    # Log average training discriminator loss\n        self.logs[\"Val_tg_loss\"].append(val_tg_loss)        # Log average validation generator loss\n        self.logs[\"Val_td_loss\"].append(val_td_loss)        # Log average validation discriminator loss\n        self.logs[\"Train_sg_loss\"].append(train_sg_loss)    # Log average training generator loss\n        self.logs[\"Train_sd_loss\"].append(train_sd_loss)    # Log average training discriminator loss\n        self.logs[\"Val_sg_loss\"].append(val_sg_loss)        # Log average validation generator loss\n        self.logs[\"Val_sd_loss\"].append(val_sd_loss)        # Log average validation discriminator loss\n        self.logs[\"Samples\"].append(make_grid(self.target_generator(self.z).cpu() * 0.5 + 0.5, normalize=True))  # Log generated sample images\n\n    def save_model(self, full: bool = False):\n        \"\"\"\n        Save the generator and discriminator models.\n        \n        Args:\n            full (bool, optional): Whether to save the full model or just the state dictionaries. Defaults to False.\n        \"\"\"\n        if full:\n            # Save the entire generator and discriminator models\n            torch.save(self.target_generator, Path(self.save_path) / \"target_generator.pth\")\n            torch.save(self.target_discriminator, Path(self.save_path) / \"target_discriminator.pth\")\n            torch.save(self.source_generator, Path(self.save_path) / \"source_generator.pth\")\n            torch.save(self.source_discriminator, Path(self.save_path) / \"source_discriminator.pth\")\n        else:\n            # Save only the state dictionaries (weights) of the models\n            torch.save(\n                self.target_generator.state_dict(),\n                Path(self.save_path) / \"target_generator_weights.pth\",\n            )\n            torch.save(\n                self.target_discriminator.state_dict(),\n                Path(self.save_path) / \"target_discriminator_weights.pth\",\n            )\n            torch.save(\n                self.source_generator.state_dict(),\n                Path(self.save_path) / \"source_generator_weights.pth\",\n            )\n            torch.save(\n                self.source_discriminator.state_dict(),\n                Path(self.save_path) / \"source_discriminator_weights.pth\",\n            )\n\n    def load_model(self, full: bool = False):\n        \"\"\"\n        Load the generator and discriminator models.\n        \n        Args:\n            full (bool, optional): Whether to load the full model or just the state dictionaries. Defaults to False.\n        \"\"\"\n        if (\n            full\n            and (self.save_path / \"target_generator.pth\").is_file()\n            and (self.save_path / \"target_discriminator.pth\").is_file()\n            and (self.save_path / \"source_generator.pth\").is_file()\n            and (self.save_path / \"source_discriminator.pth\").is_file()\n        ):\n            # Load the entire generator and discriminator models\n            self.target_generator = torch.load(self.save_path / \"target_generator.pth\", weights_only = True)\n            self.target_discriminator = torch.load(self.save_path / \"target_discriminator.pth\", weights_only = True)\n            self.source_generator = torch.load(self.save_path / \"source_generator.pth\", weights_only = True)\n            self.source_discriminator = torch.load(self.save_path / \"source_discriminator.pth\", weights_only = True)\n        elif (\n            (self.save_path / \"target_generator_weights.pth\").is_file() and\n            (self.save_path / \"target_discriminator_weights.pth\").is_file() and\n            (self.save_path / \"source_generator_weights.pth\").is_file() and\n            (self.save_path / \"source_discriminator_weights.pth\").is_file()\n        ):\n            # Load only the state dictionaries (weights) of the models\n            self.target_generator.load_state_dict(torch.load(self.save_path / \"target_generator_weights.pth\", weights_only = True))\n            self.target_discriminator.load_state_dict(torch.load(self.save_path / \"target_discriminator_weights.pth\", weights_only = True))\n            self.source_generator.load_state_dict(torch.load(self.save_path / \"source_generator_weights.pth\", weights_only = True))\n            self.source_discriminator.load_state_dict(torch.load(self.save_path / \"source_discriminator_weights.pth\", weights_only = True))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-18T08:31:59.999066Z","iopub.execute_input":"2025-01-18T08:31:59.999596Z","iopub.status.idle":"2025-01-18T08:32:00.057844Z","shell.execute_reply.started":"2025-01-18T08:31:59.999551Z","shell.execute_reply":"2025-01-18T08:32:00.057037Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!mkdir models","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-18T08:32:00.058783Z","iopub.execute_input":"2025-01-18T08:32:00.059023Z","iopub.status.idle":"2025-01-18T08:32:01.067611Z","shell.execute_reply.started":"2025-01-18T08:32:00.059000Z","shell.execute_reply":"2025-01-18T08:32:01.066586Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Set the device to GPU if available, otherwise default to CPU\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nprint(f'Using device: {device}')\n\n# Initialize the Trainer with the necessary components and parameters\ntrainer = Trainer(\n    train_data=train_dataloader,                             # DataLoader for the training dataset\n    val_data=test_dataloader,                                # DataLoader for the validation dataset\n    target_generator=Generator(in_channels=3, features=64),  # Instantiate the Generator model with 3 input channels and 64 feature maps\n    target_discriminator=Discriminator(),                    # Instantiate the Discriminator model with default parameters\n    source_generator=Generator(in_channels=3, features=64),  # Instantiate the Generator model with 3 input channels and 64 feature maps\n    source_discriminator=Discriminator(),                    # Instantiate the Discriminator model with default parameters\n    nb_epochs=NB_EPOCHS,                                     # Set the number of training epochs to 20\n    device=device,                                           # Specify the device ('cuda' or 'cpu') for training\n    save_path='./models/',                                   # Directory path where the trained models will be saved\n    discriminator_workouts_per_step= 1,\n    generator_workouts_per_step = 1,\n    lambda_coef = LAMBDA_COEF,\n    buffer_size = BATCH_SIZE,\n    load_m = False,\n)\n\ntrainer.load_model()       # Load pre-trained model weights from the specified save path, if available\ntrainer.init_optimizers()  # Initialize the optimizers for both the generator and discriminator models\ntrainer.init_buffers()\nlogs = trainer.train()     # Begin the training process and store the training logs for later analysis","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-18T08:32:01.069038Z","iopub.execute_input":"2025-01-18T08:32:01.069347Z"}},"outputs":[],"execution_count":null}]}