{"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":30716,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# **Convolutional Block for generator**","metadata":{}},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\n\n\nclass ConvolutionalBlock(nn.Module):\n    def __init__(\n        self,\n        in_channels: int,\n        out_channels: int,\n        is_downsampling: bool = True,\n        add_activation: bool = True,\n        **kwargs\n    ):\n        super().__init__()\n        if is_downsampling:\n            self.conv = nn.Sequential(\n                nn.Conv2d(in_channels, out_channels, padding_mode=\"reflect\", **kwargs),\n                nn.InstanceNorm2d(out_channels),\n                nn.ReLU(inplace=True) if add_activation else nn.Identity(),\n            )\n        else:\n            self.conv = nn.Sequential(\n                nn.ConvTranspose2d(in_channels, out_channels, **kwargs),\n                nn.InstanceNorm2d(out_channels),\n                nn.ReLU(inplace=True) if add_activation else nn.Identity(),\n            )\n\n    def forward(self, x):\n        return self.conv(x)","metadata":{"execution":{"iopub.status.busy":"2024-06-02T20:07:15.696672Z","iopub.execute_input":"2024-06-02T20:07:15.697080Z","iopub.status.idle":"2024-06-02T20:07:15.705778Z","shell.execute_reply.started":"2024-06-02T20:07:15.697045Z","shell.execute_reply":"2024-06-02T20:07:15.704720Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **Residual block as a part of the generator.**","metadata":{}},{"cell_type":"code","source":"class ResidualBlock(nn.Module):\n    def __init__(self, channels: int):\n        super().__init__()\n        self.block = nn.Sequential(\n            ConvolutionalBlock(channels, channels, add_activation=True, kernel_size=3, padding=1),\n            ConvolutionalBlock(channels, channels, add_activation=False, kernel_size=3, padding=1),\n        )\n        \n    def forward(self, x):\n        return x + self.block(x)","metadata":{"execution":{"iopub.status.busy":"2024-06-02T20:07:18.315762Z","iopub.execute_input":"2024-06-02T20:07:18.316138Z","iopub.status.idle":"2024-06-02T20:07:18.322271Z","shell.execute_reply.started":"2024-06-02T20:07:18.316107Z","shell.execute_reply":"2024-06-02T20:07:18.321393Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **Generator**","metadata":{}},{"cell_type":"code","source":"class Generator(nn.Module):\n    def __init__(\n        self, img_channels: int, num_features: int = 64, num_residuals: int = 9\n    ):\n        super().__init__()\n        self.initial_layer = 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.ReLU(inplace=True),\n        )\n\n        self.downsampling_layers = nn.ModuleList(\n            [\n                ConvolutionalBlock(\n                    num_features,\n                    num_features * 2,\n                    is_downsampling=True,\n                    kernel_size=3,\n                    stride=2,\n                    padding=1,\n                ),\n                ConvolutionalBlock(\n                    num_features * 2,\n                    num_features * 4,\n                    is_downsampling=True,\n                    kernel_size=3,\n                    stride=2,\n                    padding=1,\n                ),\n            ]\n        )\n\n        self.residual_layers = nn.Sequential(\n            *[ResidualBlock(num_features * 4) for _ in range(num_residuals)]\n        )\n\n        self.upsampling_layers = nn.ModuleList(\n            [\n                ConvolutionalBlock(\n                    num_features * 4,\n                    num_features * 2,\n                    is_downsampling=False,\n                    kernel_size=3,\n                    stride=2,\n                    padding=1,\n                    output_padding=1,\n                ),\n                ConvolutionalBlock(\n                    num_features * 2,\n                    num_features * 1,\n                    is_downsampling=False,\n                    kernel_size=3,\n                    stride=2,\n                    padding=1,\n                    output_padding=1,\n                ),\n            ]\n        )\n\n        self.last_layer = 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_layer(x)\n        for layer in self.downsampling_layers:\n            x = layer(x)\n        x = self.residual_layers(x)\n        for layer in self.upsampling_layers:\n            x = layer(x)\n        return torch.tanh(self.last_layer(x))","metadata":{"execution":{"iopub.status.busy":"2024-06-02T20:07:20.216986Z","iopub.execute_input":"2024-06-02T20:07:20.217386Z","iopub.status.idle":"2024-06-02T20:07:20.229851Z","shell.execute_reply.started":"2024-06-02T20:07:20.217349Z","shell.execute_reply":"2024-06-02T20:07:20.228797Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class ConvInstanceNormLeakyReLUBlock(nn.Module):\n    def __init__(self, in_channels: int, out_channels: int, stride: int):\n        \"\"\"\n        Class object initialization for Convolution-InstanceNorm-LeakyReLU layer\n\n        We use leaky ReLUs with a slope of 0.2.\n        \"\"\"\n        super().__init__()\n        self.conv = nn.Sequential(\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\",\n            ),\n            nn.InstanceNorm2d(out_channels),\n            nn.LeakyReLU(0.2, inplace=True),\n        )\n\n    def forward(self, x):\n        return self.conv(x)","metadata":{"execution":{"iopub.status.busy":"2024-06-02T20:07:23.152768Z","iopub.execute_input":"2024-06-02T20:07:23.153140Z","iopub.status.idle":"2024-06-02T20:07:23.159842Z","shell.execute_reply.started":"2024-06-02T20:07:23.153107Z","shell.execute_reply":"2024-06-02T20:07:23.158924Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **Discriminator**","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_layer = 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                ConvInstanceNormLeakyReLUBlock(\n                    in_channels,\n                    feature,\n                    stride=1 if feature == features[-1] else 2,\n                )\n            )\n            in_channels = feature\n\n        # After the last layer, we apply a convolution to produce a 1-dimensional output\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        self.model = nn.Sequential(*layers)\n\n    def forward(self, x):\n        x = self.initial_layer(x)\n\n        # feed the model output into a sigmoid function to make a 1/0 label\n        return torch.sigmoid(self.model(x))","metadata":{"execution":{"iopub.status.busy":"2024-06-02T20:07:25.564069Z","iopub.execute_input":"2024-06-02T20:07:25.564725Z","iopub.status.idle":"2024-06-02T20:07:25.573811Z","shell.execute_reply.started":"2024-06-02T20:07:25.564690Z","shell.execute_reply":"2024-06-02T20:07:25.572809Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **Dataset**","metadata":{}},{"cell_type":"code","source":"from PIL import Image\nimport os\nfrom torch.utils.data import Dataset\nimport numpy as np\n\nclass MonetDataset(Dataset):\n    def __init__(self, root_monet, root_photos, transform=None):\n        self.root_monet = root_monet\n        self.root_photos = root_photos\n        self.transform = transform\n\n        self.monet_images = os.listdir(root_monet)\n        self.photos_images = os.listdir(root_photos)\n        self.length_dataset = max(len(self.monet_images), len(self.photos_images))\n        self.monet_len = len(self.monet_images)\n        self.photos_len = len(self.photos_images)\n\n    def __len__(self):\n        return self.length_dataset\n\n    def __getitem__(self, index):\n        monet_img = self.monet_images[index % self.monet_len]\n        photos_img = self.photos_images[index % self.photos_len]\n        \n        monet_path = os.path.join(self.root_monet, monet_img)\n        photos_path = os.path.join(self.root_photos, photos_img)\n\n        monet_img = np.array(Image.open(monet_path).convert(\"RGB\"))\n        photos_img = np.array(Image.open(photos_path).convert(\"RGB\"))\n\n        if self.transform:\n            augmentations = self.transform(image=monet_img, image0=photos_img)\n            monet_img = augmentations[\"image\"]\n            photos_img = augmentations[\"image0\"]\n\n        return monet_img, photos_img","metadata":{"execution":{"iopub.status.busy":"2024-06-02T20:07:28.462125Z","iopub.execute_input":"2024-06-02T20:07:28.463073Z","iopub.status.idle":"2024-06-02T20:07:28.473208Z","shell.execute_reply.started":"2024-06-02T20:07:28.463037Z","shell.execute_reply":"2024-06-02T20:07:28.472298Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **Transformer to preprocess the image**","metadata":{}},{"cell_type":"code","source":"import torch\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\n\nDEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nprint(f\"using device: {DEVICE}\")\n\nTRAIN_DIR = \"/kaggle/input/gan-getting-started\"\n# VAL_DIR = \"data/val\"\n\nBATCH_SIZE = 10\nLEARNING_RATE = 3e-4\nLAMBDA_IDENTITY = 0.0 # loss weight for identity loss\nLAMBDA_CYCLE = 10\nNUM_WORKERS = 4\nNUM_EPOCHS = 5\nLOAD_MODEL = False\nSAVE_MODEL = True\nCHECKPOINT_GENERATOR_P = \"/kaggle/working/models/genh.pth.tar\"\nCHECKPOINT_GENERATOR_M = \"/kaggle/working/models/genz.pth.tar\"\nCHECKPOINT_DISCRIMINATOR_P = \"/kaggle/working/models/disch.pth.tar\"\nCHECKPOINT_DISCRIMINATOR_M = \"/kaggle/working/models/discz.pth.tar\"\n\ntransforms = A.Compose(\n    [\n        A.Resize(width=256, height=256),\n        A.HorizontalFlip(p=0.5),\n        A.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], max_pixel_value=255),\n        ToTensorV2(),\n    ],\n    additional_targets={\"image0\": \"image\"},\n)","metadata":{"execution":{"iopub.status.busy":"2024-06-02T20:07:31.739282Z","iopub.execute_input":"2024-06-02T20:07:31.740166Z","iopub.status.idle":"2024-06-02T20:07:31.748605Z","shell.execute_reply.started":"2024-06-02T20:07:31.740131Z","shell.execute_reply":"2024-06-02T20:07:31.747753Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import random, torch, os, numpy as np\nimport torch.nn as nn\nimport copy\nimport os\n\n# def save_checkpoint(model, optimizer, filename=\"models/checkpoint.pth.tar\"):\n#     print(\"=> Saving checkpoint\")\n#     checkpoint = {\n#         \"state_dict\": model.state_dict(),\n#         \"optimizer\": optimizer.state_dict(),\n#     }\n#     torch.save(checkpoint, filename)\n\ndef save_checkpoint(model, optimizer, filename):\n    os.makedirs(os.path.dirname(filename), exist_ok=True)  # Ensure directory exists\n    print(\"=> Saving checkpoint\")\n    checkpoint = {\n        \"state_dict\": model.state_dict(),\n        \"optimizer\": optimizer.state_dict(),\n    }\n    torch.save(checkpoint, filename)\n\n\n\ndef load_checkpoint(checkpoint_file, model, optimizer, lr):\n    print(\"=> Loading checkpoint\")\n    checkpoint = torch.load(checkpoint_file, map_location=DEVICE)\n    model.load_state_dict(checkpoint[\"state_dict\"])\n    optimizer.load_state_dict(checkpoint[\"optimizer\"])\n\n    for param_group in optimizer.param_groups:\n        param_group[\"lr\"] = lr\n        \ndef seed_everything(seed=34):\n    os.environ[\"PYTHONHASHSEED\"] = str(seed)\n    random.seed(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.cuda.manual_seed_all(seed)\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = False","metadata":{"execution":{"iopub.status.busy":"2024-06-02T20:07:34.983508Z","iopub.execute_input":"2024-06-02T20:07:34.984133Z","iopub.status.idle":"2024-06-02T20:07:34.993247Z","shell.execute_reply.started":"2024-06-02T20:07:34.984094Z","shell.execute_reply":"2024-06-02T20:07:34.992359Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **Training**","metadata":{}},{"cell_type":"code","source":"import torch\nimport sys\nfrom torch.utils.data import DataLoader\nimport torch.nn as nn\nimport torch.optim as optim\nfrom tqdm import tqdm\nfrom torchvision.utils import save_image\n\ndef train_fn(\n    disc_P, disc_M, gen_M, gen_P, loader, opt_disc, opt_gen, l1, mse, d_scaler, g_scaler, epoch_num\n):\n    P_reals = 0\n    P_fakes = 0\n    loop = tqdm(loader, leave=True)\n\n    for idx, (monet, photo) in enumerate(loop):\n        monet = monet.to(DEVICE)\n        photo = photo.to(DEVICE)\n\n        # Train discriminators H and Z\n        with torch.cuda.amp.autocast():\n            fake_photo = gen_P(monet)\n            D_P_real = disc_P(photo)\n            D_P_fake = disc_P(fake_photo.detach())\n            P_reals += D_P_real.mean().item()\n            P_fakes += D_P_fake.mean().item()\n            D_P_real_loss = mse(D_P_real, torch.ones_like(D_P_real))\n            D_P_fake_loss = mse(D_P_fake, torch.zeros_like(D_P_fake))\n            D_P_loss = D_P_real_loss + D_P_fake_loss\n\n            fake_monet = gen_M(monet)\n            D_M_real = disc_M(photo)\n            D_M_fake = disc_M(fake_monet.detach())\n            D_M_real_loss = mse(D_M_real, torch.ones_like(D_M_real))\n            D_M_fake_loss = mse(D_M_fake, torch.zeros_like(D_M_fake))\n            D_M_loss = D_M_real_loss + D_M_fake_loss\n\n            D_loss = (D_M_loss + D_M_loss) / 2\n\n        opt_disc.zero_grad()\n        d_scaler.scale(D_loss).backward()\n        d_scaler.step(opt_disc)\n        d_scaler.update()\n\n        # Train generators H and Z\n        with torch.cuda.amp.autocast():\n            # adversarial losses\n            D_P_fake = disc_P(fake_photo)\n            D_M_fake = disc_M(fake_monet)\n            loss_G_P = mse(D_P_fake, torch.ones_like(D_P_fake))\n            loss_G_M = mse(D_M_fake, torch.ones_like(D_M_fake))\n\n            # cycle losses\n            cycle_monet = gen_M(fake_photo)\n            cycle_photo = gen_P(fake_monet)\n            cycle_monet_loss = l1(monet, cycle_monet)\n            cycle_photo_loss = l1(photo, cycle_photo)\n\n            # total loss\n            G_loss = (\n                loss_G_M\n                + loss_G_P\n                + cycle_monet_loss * LAMBDA_CYCLE\n                + cycle_photo_loss * LAMBDA_CYCLE\n                # + identity_horse_loss * LAMBDA_IDENTITY\n                # + identity_zebra_loss * LAMBDA_IDENTITY\n            )\n\n        opt_gen.zero_grad()\n        g_scaler.scale(G_loss).backward()\n        g_scaler.step(opt_gen)\n        g_scaler.update()\n\n        if idx % 7000 == 0: # first and last of each epoch only\n            save_image(fake_monet * 0.5 + 0.5, f\"/kaggle/working/outputs/fakemonet/epoch_{epoch_num + 1}/monet_{idx}.png\")\n\n        loop.set_postfix(P_real=P_reals / (idx + 1), P_fake=P_fakes / (idx + 1))","metadata":{"execution":{"iopub.status.busy":"2024-06-02T20:07:38.060150Z","iopub.execute_input":"2024-06-02T20:07:38.060883Z","iopub.status.idle":"2024-06-02T20:07:38.075553Z","shell.execute_reply.started":"2024-06-02T20:07:38.060849Z","shell.execute_reply":"2024-06-02T20:07:38.074661Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def main():\n    disc_P = Discriminator(in_channels=3).to(DEVICE)\n    disc_M = Discriminator(in_channels=3).to(DEVICE)\n    gen_M = Generator(img_channels=3, num_residuals=9).to(DEVICE)\n    gen_P = Generator(img_channels=3, num_residuals=9).to(DEVICE)\n\n    # use Adam Optimizer for both generator and discriminator\n    opt_disc = optim.Adam(\n        list(disc_P.parameters()) + list(disc_M.parameters()),\n        lr=LEARNING_RATE,\n        betas=(0.5, 0.999),\n    )\n\n    opt_gen = optim.Adam(\n        list(gen_M.parameters()) + list(gen_P.parameters()),\n        lr=LEARNING_RATE,\n        betas=(0.5, 0.999),\n    )\n\n    L1 = nn.L1Loss()\n    mse = nn.MSELoss()\n    global LOAD_MODEL\n    if LOAD_MODEL == True:\n        load_checkpoint(\n            CHECKPOINT_GENERATOR_P,\n            gen_P,\n            opt_gen,\n            LEARNING_RATE,\n        )\n        load_checkpoint(\n            CHECKPOINT_GENERATOR_M,\n            gen_M,\n            opt_gen,\n            LEARNING_RATE,\n        )\n        load_checkpoint(\n            CHECKPOINT_DISCRIMINATOR_P,\n            disc_M,\n            opt_disc,\n            LEARNING_RATE,\n        )\n        load_checkpoint(\n            CHECKPOINT_DISCRIMINATOR_M,\n            disc_P,\n            opt_disc,\n            LEARNING_RATE,\n        )\n    \n    dataset = MonetDataset(\n        root_photos=TRAIN_DIR + \"/photo_jpg\",\n        root_monet=TRAIN_DIR + \"/monet_jpg\",\n        transform=transforms,\n    )\n\n    loader = DataLoader(\n        dataset,\n        batch_size=BATCH_SIZE,\n        shuffle=True,\n        num_workers=NUM_WORKERS,\n        pin_memory=True,\n    )\n    g_scaler = torch.cuda.amp.GradScaler()\n    d_scaler = torch.cuda.amp.GradScaler()\n\n    for epoch in range(NUM_EPOCHS):\n        os.makedirs(f'/kaggle/working/outputs/fakephotos/epoch_{epoch + 1}/', exist_ok=True)\n        os.makedirs(f'/kaggle/working/outputs/fakemonet/epoch_{epoch + 1}/', exist_ok=True)\n        train_fn(\n            disc_P,\n            disc_M,\n            gen_M,\n            gen_P,\n            loader,\n            opt_disc,\n            opt_gen,\n            L1,\n            mse,\n            d_scaler,\n            g_scaler,\n            epoch,\n        )\n\n        if SAVE_MODEL:\n            save_checkpoint(gen_P, opt_gen, filename=CHECKPOINT_GENERATOR_P)\n            save_checkpoint(gen_M, opt_gen, filename=CHECKPOINT_GENERATOR_M)\n            save_checkpoint(disc_P, opt_disc, filename=CHECKPOINT_DISCRIMINATOR_P)\n            save_checkpoint(disc_M, opt_disc, filename=CHECKPOINT_DISCRIMINATOR_M)\n            LOAD_MODEL = True","metadata":{"execution":{"iopub.status.busy":"2024-06-02T20:07:41.235585Z","iopub.execute_input":"2024-06-02T20:07:41.236235Z","iopub.status.idle":"2024-06-02T20:07:41.248878Z","shell.execute_reply.started":"2024-06-02T20:07:41.236199Z","shell.execute_reply":"2024-06-02T20:07:41.247968Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"main()","metadata":{"execution":{"iopub.status.busy":"2024-06-02T20:07:59.360782Z","iopub.execute_input":"2024-06-02T20:07:59.361548Z","iopub.status.idle":"2024-06-02T21:41:54.103894Z","shell.execute_reply.started":"2024-06-02T20:07:59.361511Z","shell.execute_reply":"2024-06-02T21:41:54.102763Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nos.makedirs('/kaggle/working/outputs/fakephotos/', exist_ok=True)\nos.makedirs('/kaggle/working/outputs/fakemonet/', exist_ok=True)\n# os.makedirs('/kaggle/working/models/', exist_ok = True)","metadata":{"execution":{"iopub.status.busy":"2024-06-02T20:07:52.026166Z","iopub.execute_input":"2024-06-02T20:07:52.026874Z","iopub.status.idle":"2024-06-02T20:07:52.031741Z","shell.execute_reply.started":"2024-06-02T20:07:52.026842Z","shell.execute_reply":"2024-06-02T20:07:52.030753Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from shutil import rmtree\nrmtree('/kaggle/working/outputs/fakephotos/')\nrmtree('/kaggle/working/outputs/fakemonet/')","metadata":{"execution":{"iopub.status.busy":"2024-06-02T21:44:00.273831Z","iopub.execute_input":"2024-06-02T21:44:00.274760Z","iopub.status.idle":"2024-06-02T21:44:00.282472Z","shell.execute_reply.started":"2024-06-02T21:44:00.274720Z","shell.execute_reply":"2024-06-02T21:44:00.281672Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def main():\n    disc_P = Discriminator(in_channels=3).to(DEVICE)\n    disc_M = Discriminator(in_channels=3).to(DEVICE)\n    gen_M = Generator(img_channels=3, num_residuals=9).to(DEVICE)\n    gen_P = Generator(img_channels=3, num_residuals=9).to(DEVICE)\n\n    # use Adam Optimizer for both generator and discriminator\n    opt_disc = optim.Adam(\n        list(disc_P.parameters()) + list(disc_M.parameters()),\n        lr=LEARNING_RATE,\n        betas=(0.5, 0.999),\n    )\n\n    opt_gen = optim.Adam(\n        list(gen_M.parameters()) + list(gen_P.parameters()),\n        lr=LEARNING_RATE,\n        betas=(0.5, 0.999),\n    )\n\n    L1 = nn.L1Loss()\n    mse = nn.MSELoss()\n    global LOAD_MODEL\n    if LOAD_MODEL == True:\n        load_checkpoint(\n            CHECKPOINT_GENERATOR_P,\n            gen_P,\n            opt_gen,\n            LEARNING_RATE,\n        )\n        load_checkpoint(\n            CHECKPOINT_GENERATOR_M,\n            gen_M,\n            opt_gen,\n            LEARNING_RATE,\n        )\n        load_checkpoint(\n            CHECKPOINT_DISCRIMINATOR_P,\n            disc_M,\n            opt_disc,\n            LEARNING_RATE,\n        )\n        load_checkpoint(\n            CHECKPOINT_DISCRIMINATOR_M,\n            disc_P,\n            opt_disc,\n            LEARNING_RATE,\n        )\n\n    dataset = MonetDataset(\n        root_photos=TRAIN_DIR + \"/photo_jpg\",\n        root_monet=TRAIN_DIR + \"/monet_jpg\",\n        transform=transforms,\n    )\n    # val_dataset = MonetDataset(\n    #     root_photos=VAL_DIR + \"/horses\",\n    #     root_monet=VAL_DIR + \"/zebras\",\n    #     transform=transforms,\n    # )\n    # val_loader = DataLoader(\n    #     val_dataset,\n    #     batch_size=BATCH_SIZE,\n    #     shuffle=False,\n    #     pin_memory=True,\n    # )\n    \n    loader = DataLoader(\n        dataset,\n        batch_size=BATCH_SIZE,\n        shuffle=True,\n        num_workers=NUM_WORKERS,\n        pin_memory=True,\n    )\n    g_scaler = torch.cuda.amp.GradScaler()\n    d_scaler = torch.cuda.amp.GradScaler()\n\n    os.makedirs(f'/kaggle/working/testout/fakemonet/epoch_{epoch + 1}/', exist_ok=True)\n    train_fn(\n        disc_P,\n        disc_M,\n        gen_M,\n        gen_P,\n        loader,\n        opt_disc,\n        opt_gen,\n        L1,\n        mse,\n        d_scaler,\n        g_scaler,\n        epoch,\n    )\n    ","metadata":{"execution":{"iopub.status.busy":"2024-06-02T21:45:05.688083Z","iopub.execute_input":"2024-06-02T21:45:05.688811Z","iopub.status.idle":"2024-06-02T21:45:05.700566Z","shell.execute_reply.started":"2024-06-02T21:45:05.688776Z","shell.execute_reply":"2024-06-02T21:45:05.699582Z"},"trusted":true},"execution_count":null,"outputs":[]}]}