{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","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":"nvidiaTeslaT4","dataSources":[{"sourceId":21755,"databundleVersionId":1475600,"sourceType":"competition"},{"sourceId":1217826,"sourceType":"datasetVersion","datasetId":298806},{"sourceId":1504266,"sourceType":"datasetVersion","datasetId":885385},{"sourceId":2520560,"sourceType":"datasetVersion","datasetId":1526975},{"sourceId":3956508,"sourceType":"datasetVersion","datasetId":850761},{"sourceId":7457578,"sourceType":"datasetVersion","datasetId":1698586}],"dockerImageVersionId":30787,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import torch\n\nimport torch.nn as nn\n\nfrom PIL import Image\nimport os\nimport numpy as np\nimport torch\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nimport sys\nfrom tqdm import tqdm\nfrom torch.utils.data import DataLoader\nimport torch.optim as optim\nfrom torchvision.utils import save_image\nfrom torch.utils.data import Dataset\nimport numpy as np\nimport warnings\nwarnings.filterwarnings(\"ignore\")","metadata":{"execution":{"iopub.status.busy":"2024-11-09T06:19:51.879134Z","iopub.execute_input":"2024-11-09T06:19:51.879475Z","iopub.status.idle":"2024-11-09T06:19:57.865864Z","shell.execute_reply.started":"2024-11-09T06:19:51.879442Z","shell.execute_reply":"2024-11-09T06:19:57.865088Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"base_path = '/kaggle/working/' \nmonet_path = os.path.join(base_path, 'monet') \nactual_path = os.path.join(base_path, 'actual')\nos.makedirs(monet_path, exist_ok=True)\nos.makedirs(actual_path, exist_ok=True)","metadata":{"execution":{"iopub.status.busy":"2024-11-09T06:19:57.867257Z","iopub.execute_input":"2024-11-09T06:19:57.867655Z","iopub.status.idle":"2024-11-09T06:19:57.873346Z","shell.execute_reply.started":"2024-11-09T06:19:57.867620Z","shell.execute_reply":"2024-11-09T06:19:57.872326Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Define the new dataset paths\n\nHUMAN_PATH = \"/kaggle/input/human-faces/Humans\"\nPORTRAIT_PATHS = [\n    \"/kaggle/input/portrait-paintings/Images\",\n    \"/kaggle/input/art-portraits/Portraits/Portraits\",\n    \"/kaggle/input/art-portraits/Portraits_update/Portraits\"\n]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-09T06:19:57.874604Z","iopub.execute_input":"2024-11-09T06:19:57.875259Z","iopub.status.idle":"2024-11-09T06:19:57.881075Z","shell.execute_reply.started":"2024-11-09T06:19:57.875212Z","shell.execute_reply":"2024-11-09T06:19:57.880367Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# **Configuration**","metadata":{}},{"cell_type":"code","source":"def save_checkpoint(model, optimizer, filename=\"/kaggle/working/.pth.tar\"):\n\n    print(\"=> Saving checkpoint\")\n\n    checkpoint = {\n\n        \"state_dict\": model.state_dict(),\n\n        \"optimizer\": optimizer.state_dict(),\n\n    }\n\n    torch.save(checkpoint, filename)\n\n\n\n\n\ndef load_checkpoint(checkpoint_file, model, optimizer, lr):\n\n    print(\"=> Loading checkpoint\")\n\n    checkpoint = torch.load(checkpoint_file, map_location=DEVICE)\n\n    model.load_state_dict(checkpoint[\"state_dict\"])\n\n    optimizer.load_state_dict(checkpoint[\"optimizer\"])\n\n\n\n    # If we don't do this then it will just have learning rate of old checkpoint\n\n    # and it will lead to many hours of debugging \\:\n\n    for param_group in optimizer.param_groups:\n\n        param_group[\"lr\"] = lr\n\n\n\n\n\ndef seed_everything(seed=42):\n\n    os.environ[\"PYTHONHASHSEED\"] = str(seed)\n\n    random.seed(seed)\n\n    np.random.seed(seed)\n\n    torch.manual_seed(seed)\n\n    torch.cuda.manual_seed(seed)\n\n    torch.cuda.manual_seed_all(seed)\n\n    torch.backends.cudnn.deterministic = True\n\n    torch.backends.cudnn.benchmark = False\n\n\n\nDEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\n#TRAIN_DIR = \"data/train\"\n\n#VAL_DIR = \"data/val\"\n\nBATCH_SIZE = 20\n\nLEARNING_RATE = 1e-5\n\nLAMBDA_IDENTITY = 0.0\n\nLAMBDA_CYCLE = 10\n\nNUM_WORKERS = 4\n\nNUM_EPOCHS = 19\n\nLOAD_MODEL = False\n\nSAVE_MODEL = True\n\nCHECKPOINT_GEN_A = \"/kaggle/working/genH.pth.tar\"\n\nCHECKPOINT_GEN_M = \"/kaggle/working/genP.pth.tar\"\n\nCHECKPOINT_CRITIC_A = \"/kaggle/working/criticH.pth.tar\"\n\nCHECKPOINT_CRITIC_M = \"/kaggle/working/criticP.pth.tar\"\n\n\n\ntransforms = A.Compose(\n\n    [\n\n        A.Resize(width=256, height=256),\n\n        A.HorizontalFlip(p=0.5),\n\n        A.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], max_pixel_value=255),\n\n        ToTensorV2(),\n\n    ],\n\n    additional_targets={\"image0\": \"image\"},\n    is_check_shapes=False\n\n)","metadata":{"execution":{"iopub.status.busy":"2024-11-09T14:28:30.101676Z","iopub.execute_input":"2024-11-09T14:28:30.102064Z","iopub.status.idle":"2024-11-09T14:28:30.115275Z","shell.execute_reply.started":"2024-11-09T14:28:30.102026Z","shell.execute_reply":"2024-11-09T14:28:30.114324Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# **Generator**","metadata":{}},{"cell_type":"code","source":"class ConvBlock(nn.Module):\n\n    def __init__(self, in_channels, out_channels, down=True, use_act=True, **kwargs):\n\n        super().__init__()\n\n        self.conv = nn.Sequential(\n\n            nn.Conv2d(in_channels, out_channels, padding_mode=\"reflect\", **kwargs)\n\n            if down\n\n            else nn.ConvTranspose2d(in_channels, out_channels, **kwargs),\n\n            nn.InstanceNorm2d(out_channels),\n\n            nn.ReLU(inplace=True) if use_act else nn.Identity(),\n\n        )\n\n\n\n    def forward(self, x):\n\n        return self.conv(x)\n\n\n\n\n\nclass ResidualBlock(nn.Module):\n\n    def __init__(self, channels):\n\n        super().__init__()\n\n        self.block = nn.Sequential(\n\n            ConvBlock(channels, channels, kernel_size=3, padding=1),\n\n            ConvBlock(channels, channels, use_act=False, kernel_size=3, padding=1),\n\n        )\n\n\n\n    def forward(self, x):\n\n        return x + self.block(x)\n\n\n\n\n\nclass Generator(nn.Module):\n\n    def __init__(self, img_channels, num_features=64, num_residuals=9):\n\n        super().__init__()\n\n        self.initial = nn.Sequential(\n\n            nn.Conv2d(\n\n                img_channels,\n\n                num_features,\n\n                kernel_size=7,\n\n                stride=1,\n\n                padding=3,\n\n                padding_mode=\"reflect\",\n\n            ),\n\n            nn.InstanceNorm2d(num_features),\n\n            nn.ReLU(inplace=True),\n\n        )\n\n        self.down_blocks = nn.ModuleList(\n\n            [\n\n                ConvBlock(\n\n                    num_features, num_features * 2, kernel_size=3, stride=2, padding=1\n\n                ),\n\n                ConvBlock(\n\n                    num_features * 2,\n\n                    num_features * 4,\n\n                    kernel_size=3,\n\n                    stride=2,\n\n                    padding=1,\n\n                ),\n\n            ]\n\n        )\n\n        self.res_blocks = nn.Sequential(\n\n            *[ResidualBlock(num_features * 4) for _ in range(num_residuals)]\n\n        )\n\n        self.up_blocks = nn.ModuleList(\n\n            [\n\n                ConvBlock(\n\n                    num_features * 4,\n\n                    num_features * 2,\n\n                    down=False,\n\n                    kernel_size=3,\n\n                    stride=2,\n\n                    padding=1,\n\n                    output_padding=1,\n\n                ),\n\n                ConvBlock(\n\n                    num_features * 2,\n\n                    num_features * 1,\n\n                    down=False,\n\n                    kernel_size=3,\n\n                    stride=2,\n\n                    padding=1,\n\n                    output_padding=1,\n\n                ),\n\n            ]\n\n        )\n\n\n\n        self.last = nn.Conv2d(\n\n            num_features * 1,\n\n            img_channels,\n\n            kernel_size=7,\n\n            stride=1,\n\n            padding=3,\n\n            padding_mode=\"reflect\",\n\n        )\n\n\n\n    def forward(self, x):\n\n        x = self.initial(x)\n\n        for layer in self.down_blocks:\n\n            x = layer(x)\n\n        x = self.res_blocks(x)\n\n        for layer in self.up_blocks:\n\n            x = layer(x)\n\n        return torch.tanh(self.last(x))\n","metadata":{"execution":{"iopub.status.busy":"2024-11-09T06:19:57.957247Z","iopub.execute_input":"2024-11-09T06:19:57.957543Z","iopub.status.idle":"2024-11-09T06:19:57.974952Z","shell.execute_reply.started":"2024-11-09T06:19:57.957511Z","shell.execute_reply":"2024-11-09T06:19:57.974076Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# **Discriminator**","metadata":{}},{"cell_type":"code","source":"class Block(nn.Module):\n\n    def __init__(self, in_channels, out_channels, stride):\n\n        super().__init__()\n\n        self.conv = nn.Sequential(\n\n            nn.Conv2d(\n\n                in_channels,\n\n                out_channels,\n\n                4,\n\n                stride,\n\n                1,\n\n                bias=True,\n\n                padding_mode=\"reflect\",\n\n            ),\n\n            nn.InstanceNorm2d(out_channels),\n\n            nn.LeakyReLU(0.2, inplace=True),\n\n        )\n\n\n\n    def forward(self, x):\n\n        return self.conv(x)\n\n\n\n\n\nclass Discriminator(nn.Module):\n\n    def __init__(self, in_channels=3, features=[64, 128, 256, 512]):\n\n        super().__init__()\n\n        self.initial = nn.Sequential(\n\n            nn.Conv2d(\n\n                in_channels,\n\n                features[0],\n\n                kernel_size=4,\n\n                stride=2,\n\n                padding=1,\n\n                padding_mode=\"reflect\",\n\n            ),\n\n            nn.LeakyReLU(0.2, inplace=True),\n\n        )\n\n\n\n        layers = []\n\n        in_channels = features[0]\n\n        for feature in features[1:]:\n\n            layers.append(\n\n                Block(in_channels, feature, stride=1 if feature == features[-1] else 2)\n\n            )\n\n            in_channels = feature\n\n        layers.append(\n\n            nn.Conv2d(\n\n                in_channels,\n\n                1,\n\n                kernel_size=4,\n\n                stride=1,\n\n                padding=1,\n\n                padding_mode=\"reflect\",\n\n            )\n\n        )\n\n        self.model = nn.Sequential(*layers)\n\n\n\n    def forward(self, x):\n\n        x = self.initial(x)\n\n        return torch.sigmoid(self.model(x))\n","metadata":{"execution":{"iopub.status.busy":"2024-11-09T06:19:57.976051Z","iopub.execute_input":"2024-11-09T06:19:57.976454Z","iopub.status.idle":"2024-11-09T06:19:57.988328Z","shell.execute_reply.started":"2024-11-09T06:19:57.976421Z","shell.execute_reply":"2024-11-09T06:19:57.987418Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# **Buliding**","metadata":{}},{"cell_type":"code","source":"class HumanToPortrait(Dataset):\n\n    def __init__(self, human_path, portrait_paths, transform=None):\n\n        self.human_path = human_path\n\n        self.portrait_paths = portrait_paths\n\n        self.transform = transform\n\n\n\n        self.human_images = os.listdir(human_path)\n\n        self.portrait_images = []\n\n        for path in portrait_paths:\n\n            self.portrait_images.extend(os.listdir(path))\n\n\n\n        self.length_dataset = max(len(self.human_images), len(self.portrait_images))\n\n        self.human_len = len(self.human_images)\n\n        self.portrait_len = len(self.portrait_images)\n\n\n\n    def __len__(self):\n\n        return self.length_dataset\n\n\n\n    def __getitem__(self, index):\n\n        human_img = self.human_images[index % self.human_len]\n\n        portrait_img = self.portrait_images[index % self.portrait_len]\n\n\n\n        human_img = np.array(Image.open(os.path.join(self.human_path, human_img)).convert(\"RGB\"))\n\n        portrait_img_path = self.portrait_images[index % self.portrait_len]\n\n        for path in self.portrait_paths:\n\n            if portrait_img_path in os.listdir(path):\n\n                portrait_img = np.array(Image.open(os.path.join(path, portrait_img_path)).convert(\"RGB\"))\n\n                break\n\n\n\n        if self.transform:\n\n            augmentations = self.transform(image=human_img, image0=portrait_img)\n\n            human_img = augmentations[\"image\"]\n\n            portrait_img = augmentations[\"image0\"]\n\n\n\n        return human_img, portrait_img","metadata":{"execution":{"iopub.status.busy":"2024-11-09T06:19:57.989485Z","iopub.execute_input":"2024-11-09T06:19:57.989783Z","iopub.status.idle":"2024-11-09T06:19:58.000288Z","shell.execute_reply.started":"2024-11-09T06:19:57.989740Z","shell.execute_reply":"2024-11-09T06:19:57.999480Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# **Loading Data**","metadata":{}},{"cell_type":"code","source":"# Create the dataset object\n\ndataset = HumanToPortrait(\n\n    human_path=HUMAN_PATH,\n\n    portrait_paths=PORTRAIT_PATHS,\n\n    transform=transforms\n\n)\n\n\n\n# Create the DataLoader\n\nloader = DataLoader(\n\n    dataset,\n\n    batch_size=BATCH_SIZE,\n\n    shuffle=True,\n\n    num_workers=NUM_WORKERS,\n\n    pin_memory=True\n\n)","metadata":{"execution":{"iopub.status.busy":"2024-11-09T06:19:58.001333Z","iopub.execute_input":"2024-11-09T06:19:58.001631Z","iopub.status.idle":"2024-11-09T06:19:58.921166Z","shell.execute_reply.started":"2024-11-09T06:19:58.001600Z","shell.execute_reply":"2024-11-09T06:19:58.920177Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# **Training Function**","metadata":{}},{"cell_type":"code","source":"def train_fn(\n    disc_H, disc_P, gen_P, gen_H, loader, opt_disc, opt_gen, l1, mse, d_scaler, g_scaler, epoch\n):\n    H_reals = 0\n    H_fakes = 0\n    P_reals = 0\n    P_fakes = 0\n    loop = tqdm(loader, leave=True)\n\n    for idx, (human, portrait) in enumerate(loop):\n\n        human = human.to(DEVICE)\n        portrait = portrait.to(DEVICE)\n\n        # Train Discriminators H and P\n        with torch.cuda.amp.autocast():\n            fake_portrait = gen_P(human)\n            D_P_real = disc_P(portrait)\n            D_P_fake = disc_P(fake_portrait.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_human = gen_H(portrait)\n            D_H_real = disc_H(human)\n            D_H_fake = disc_H(fake_human.detach())\n            D_H_real_loss = mse(D_H_real, torch.ones_like(D_H_real))\n            D_H_fake_loss = mse(D_H_fake, torch.zeros_like(D_H_fake))\n            D_H_loss = D_H_real_loss + D_H_fake_loss\n\n            # put it together\n            D_loss = (D_P_loss + D_H_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 P\n        with torch.cuda.amp.autocast():\n            # adversarial loss for both generators\n            D_P_fake = disc_P(fake_portrait)\n            D_H_fake = disc_H(fake_human)\n            loss_G_P = mse(D_P_fake, torch.ones_like(D_P_fake))\n            loss_G_H = mse(D_H_fake, torch.ones_like(D_H_fake))\n\n            # cycle loss\n            cycle_human = gen_H(fake_portrait)\n            cycle_portrait = gen_P(fake_human)\n            cycle_human_loss = l1(human, cycle_human)\n            cycle_portrait_loss = l1(portrait, cycle_portrait)\n\n            # identity loss (remove these for efficiency if you set lambda_identity=0)\n            identity_human = gen_H(human)\n            identity_portrait = gen_P(portrait)\n            identity_human_loss = l1(human, identity_human)\n            identity_portrait_loss = l1(portrait, identity_portrait)\n\n            # add all together\n            G_loss = (\n                loss_G_H\n                + loss_G_P\n                + cycle_human_loss * LAMBDA_CYCLE\n                + cycle_portrait_loss * LAMBDA_CYCLE\n                + identity_human_loss * LAMBDA_IDENTITY\n                + identity_portrait_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 % 500 == 0:\n            save_image(fake_portrait * 0.5 + 0.5, f\"/kaggle/working/portrait/FAKE_PORTRAIT_index_{idx}_epoch_{epoch}.png\")\n            save_image(fake_human * 0.5 + 0.5, f\"/kaggle/working/human/FAKE_HUMAN_index_{idx}_epoch_{epoch}.png\")\n\n        loop.set_postfix(P_real=P_reals / (idx + 1), P_fake=P_fakes / (idx + 1))","metadata":{"execution":{"iopub.status.busy":"2024-11-09T06:19:58.922547Z","iopub.execute_input":"2024-11-09T06:19:58.923319Z","iopub.status.idle":"2024-11-09T06:19:58.940111Z","shell.execute_reply.started":"2024-11-09T06:19:58.923273Z","shell.execute_reply":"2024-11-09T06:19:58.939076Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import matplotlib.pyplot as plt\n\ndef display_image(image_tensor):\n    image = image_tensor.cpu().detach().numpy().transpose(1, 2, 0)\n    image = (image * 0.5) + 0.5  # Unnormalize\n    plt.imshow(image)\n    plt.axis('off')\n    plt.show()\n\ndef main():\n    disc_H = Discriminator(in_channels=3).to(DEVICE)\n    disc_P = Discriminator(in_channels=3).to(DEVICE)\n    gen_P = Generator(img_channels=3, num_residuals=9).to(DEVICE)\n    gen_H = Generator(img_channels=3, num_residuals=9).to(DEVICE)\n\n    # Wrap models with DataParallel\n    disc_H = torch.nn.DataParallel(disc_H)\n    disc_P = torch.nn.DataParallel(disc_P)\n    gen_P = torch.nn.DataParallel(gen_P)\n    gen_H = torch.nn.DataParallel(gen_H)\n\n    opt_disc = optim.Adam(\n        list(disc_H.parameters()) + list(disc_P.parameters()),\n        lr=LEARNING_RATE,\n        betas=(0.5, 0.999),\n    )\n\n    opt_gen = optim.Adam(\n        list(gen_P.parameters()) + list(gen_H.parameters()),\n        lr=LEARNING_RATE,\n        betas=(0.5, 0.999),\n    )\n\n    L1 = nn.L1Loss()\n    mse = nn.MSELoss()\n\n    dataset = HumanToPortrait(\n        human_path=HUMAN_PATH,\n        portrait_paths=PORTRAIT_PATHS,\n        transform=transforms,\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    # Create directories if they don't exist\n    os.makedirs(\"/kaggle/working/portrait\", exist_ok=True)\n    os.makedirs(\"/kaggle/working/human\", exist_ok=True)\n\n    for epoch in range(NUM_EPOCHS):\n        train_fn(\n            disc_H,\n            disc_P,\n            gen_P,\n            gen_H,\n            loader,\n            opt_disc,\n            opt_gen,\n            L1,\n            mse,\n            d_scaler,\n            g_scaler,\n            epoch\n        )\n        save_checkpoint(gen_H, opt_gen, filename=f\"/kaggle/working/gen_H_epoch_{epoch}.pth\")\n        save_checkpoint(gen_P, opt_gen, filename=f\"/kaggle/working/gen_P_epoch_{epoch}.pth\")\n        save_checkpoint(disc_H, opt_disc, filename=f\"/kaggle/working/disc_H_epoch_{epoch}.pth\")\n        save_checkpoint(disc_P, opt_disc, filename=f\"/kaggle/working/disc_P_epoch_{epoch}.pth\")\n\n        # Generate and display an image\n        with torch.no_grad():\n            random_human_image = dataset[0][0].unsqueeze(0).to(DEVICE)  # Take the first image from the dataset\n            fake_portrait = gen_P(random_human_image)\n            display_image(fake_portrait[0])\n\nmain()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-09T06:19:58.943332Z","iopub.execute_input":"2024-11-09T06:19:58.943652Z","iopub.status.idle":"2024-11-09T14:20:25.307517Z","shell.execute_reply.started":"2024-11-09T06:19:58.943619Z","shell.execute_reply":"2024-11-09T14:20:25.305948Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import matplotlib.pyplot as plt\nimport torch\nimport torch.optim as optim\nimport torch.nn as nn\nfrom torch.utils.data import DataLoader\nimport os\n\ndef display_image(image_tensor):\n    image = image_tensor.cpu().detach().numpy().transpose(1, 2, 0)\n    image = (image * 0.5) + 0.5  # Unnormalize\n    plt.imshow(image)\n    plt.axis('off')\n    plt.show()\n\ndef load_checkpoint(checkpoint_file, model, optimizer, lr):\n    print(f\"=> Loading checkpoint from {checkpoint_file}\")\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    # If you want to change the learning rate\n    for param_group in optimizer.param_groups:\n        param_group[\"lr\"] = lr\n\ndef main():\n    print(\"Initializing models and optimizers...\")\n    disc_H = Discriminator(in_channels=3).to(DEVICE)\n    disc_P = Discriminator(in_channels=3).to(DEVICE)\n    gen_P = Generator(img_channels=3, num_residuals=9).to(DEVICE)\n    gen_H = Generator(img_channels=3, num_residuals=9).to(DEVICE)\n\n    # Wrap models with DataParallel\n    disc_H = torch.nn.DataParallel(disc_H)\n    disc_P = torch.nn.DataParallel(disc_P)\n    gen_P = torch.nn.DataParallel(gen_P)\n    gen_H = torch.nn.DataParallel(gen_H)\n\n    opt_disc = optim.Adam(\n        list(disc_H.parameters()) + list(disc_P.parameters()),\n        lr=LEARNING_RATE,\n        betas=(0.5, 0.999),\n    )\n\n    opt_gen = optim.Adam(\n        list(gen_P.parameters()) + list(gen_H.parameters()),\n        lr=LEARNING_RATE,\n        betas=(0.5, 0.999),\n    )\n\n    L1 = nn.L1Loss()\n    mse = nn.MSELoss()\n\n    print(\"Loading dataset...\")\n    dataset = HumanToPortrait(\n        human_path=HUMAN_PATH,\n        portrait_paths=PORTRAIT_PATHS,\n        transform=transforms,\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    # Load checkpoints if available\n    if os.path.exists(\"gen_H_epoch_16.pth\"):\n        load_checkpoint(\"gen_H_epoch_16.pth\", gen_H, opt_gen, LEARNING_RATE)\n    if os.path.exists(\"gen_P_epoch_16.pth\"):\n        load_checkpoint(\"gen_P_epoch_16.pth\", gen_P, opt_gen, LEARNING_RATE)\n    if os.path.exists(\"disc_H_epoch_16.pth\"):\n        load_checkpoint(\"disc_H_epoch_16.pth\", disc_H, opt_disc, LEARNING_RATE)\n    if os.path.exists(\"disc_P_epoch_16.pth\"):\n        load_checkpoint(\"disc_P_epoch_16.pth\", disc_P, opt_disc, LEARNING_RATE)\n\n    # Create directories if they don't exist\n    os.makedirs(\"/kaggle/working/portrait\", exist_ok=True)\n    os.makedirs(\"/kaggle/working/human\", exist_ok=True)\n\n    print(\"Starting training loop...\")\n    for epoch in range(17, NUM_EPOCHS):  # Start from epoch 17\n        print(f\"Epoch {epoch}/{NUM_EPOCHS}\")\n        train_fn(\n            disc_H,\n            disc_P,\n            gen_P,\n            gen_H,\n            loader,\n            opt_disc,\n            opt_gen,\n            L1,\n            mse,\n            d_scaler,\n            g_scaler,\n            epoch\n        )\n        print(f\"Saving checkpoints for epoch {epoch}...\")\n        save_checkpoint(gen_H, opt_gen, filename=f\"/kaggle/working/gen_H_epoch_{epoch}.pth\")\n        save_checkpoint(gen_P, opt_gen, filename=f\"/kaggle/working/gen_P_epoch_{epoch}.pth\")\n        save_checkpoint(disc_H, opt_disc, filename=f\"/kaggle/working/disc_H_epoch_{epoch}.pth\")\n        save_checkpoint(disc_P, opt_disc, filename=f\"/kaggle/working/disc_P_epoch_{epoch}.pth\")\n\n        # Generate and display an image\n        with torch.no_grad():\n            random_human_image = dataset[0][0].unsqueeze(0).to(DEVICE)  # Take the first image from the dataset\n            fake_portrait = gen_P(random_human_image)\n            display_image(fake_portrait[0])\n\n    print(\"Training completed.\")\n\nmain()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-09T14:28:38.443265Z","iopub.execute_input":"2024-11-09T14:28:38.443800Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}