{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.12","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"},{"sourceId":1388667,"sourceType":"datasetVersion","datasetId":810652},{"sourceId":1567182,"sourceType":"datasetVersion","datasetId":925704}],"dockerImageVersionId":30919,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"%%writefile dataset_turbo.py\n\nimport os\nimport random\nimport torch\nimport cv2\nimport numpy as np\nimport matplotlib.pyplot as plt\nfrom PIL import Image\nfrom torch.utils.data import Dataset, DataLoader\nfrom torchvision.transforms import ToTensor, Compose, Normalize\n\n\nSTYLE_FOLDERS = [\n    '/kaggle/input/vangogh2photo/vangogh2photo/trainA',\n    '/kaggle/input/monet2photo/trainA'\n]\n\nstyle_label_map = {i+1: os.path.basename(path) for i, path in enumerate(STYLE_FOLDERS)}\nprint(\"Style Label Map:\", style_label_map)\n\ndef get_all_images(style_folders):\n    \"\"\"\n    Duyệt toàn bộ ảnh trong các thư mục style và gán nhãn theo danh sách folder.\n    \"\"\"\n    image_paths = []\n    labels = []\n\n    for label_id, folder in enumerate(style_folders):\n        for root, _, files in os.walk(folder):\n            for file in files:\n                if file.lower().endswith(('png', 'jpg', 'jpeg')):\n                    image_paths.append(os.path.join(root, file))\n                    labels.append(label_id)\n\n    return image_paths, labels\n\n\nclass MultiPathStyleDataset(Dataset):\n    def __init__(self, raw_root, style_folders, raw_transform=None, style_transform=None):\n        self.raw_images_path = [os.path.join(raw_root, i) for i in os.listdir(raw_root) if\n                                i.endswith(('png', 'jpg', 'jpeg'))]\n\n        # 🔍 Quét tất cả ảnh trong nhiều thư mục style\n        self.style_paths, self.style_labels = get_all_images(style_folders)\n\n        self.raw_transform = raw_transform\n        self.style_transform = style_transform\n\n    def __len__(self):\n        return len(self.raw_images_path)\n\n    def __getitem__(self, index):\n        raw_image = cv2.imread(self.raw_images_path[index])\n        raw_image = cv2.cvtColor(raw_image, cv2.COLOR_BGR2RGB)\n        raw_image = Image.fromarray(raw_image)\n\n        style_idx = random.randint(0, len(self.style_paths) - 1)\n        style_image = cv2.imread(self.style_paths[style_idx])\n        style_image = cv2.cvtColor(style_image, cv2.COLOR_BGR2RGB)\n        style_image = Image.fromarray(style_image)\n        domain_label = self.style_labels[style_idx]\n\n        # Apply transform\n        if self.raw_transform:\n            raw_image = self.raw_transform(raw_image)\n        if self.style_transform:\n            style_image = self.style_transform(style_image)\n\n        return raw_image, style_image, torch.tensor(domain_label)\n\n\n# 📌 **Tạo dataset**\nif __name__ == '__main__':\n    STYLE_FOLDERS = [\n    '/kaggle/input/vangogh2photo/vangogh2photo/trainA',\n    '/kaggle/input/monet2photo/trainA'\n]\n    \n    dataset = MultiPathStyleDataset(\n        raw_root=\"C:/Users/LENOVO/Downloads/gan-getting-started/photo_jpg\",\n        style_folders=STYLE_FOLDERS,\n        raw_transform=Compose([\n            ToTensor(),\n            Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),\n        ]),\n        style_transform=Compose([\n            ToTensor(),\n            Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),\n        ])\n    )\n\n    data_loader = DataLoader(dataset=dataset, batch_size=8, shuffle=True, num_workers=4)\n\n    for raw_img, style_img, labels in data_loader:\n        fig, ax = plt.subplots(2, 8, figsize=(15, 6))\n        ax = ax.flatten()\n\n        for i in range(8):\n            raw_img_vis = (raw_img[i].permute(1, 2, 0) + 1) / 2\n            style_img_vis = (style_img[i].permute(1, 2, 0) + 1) / 2\n\n            ax[i * 2].imshow(raw_img_vis.cpu().numpy())\n            ax[i * 2].set_title(\"Original\")\n\n            ax[i * 2 + 1].imshow(style_img_vis.cpu().numpy())\n            ax[i * 2 + 1].set_title(style_label_map[labels[i].item()])\n\n            ax[i * 2].axis('off')\n            ax[i * 2 + 1].axis('off')\n\n        plt.show()\n        print(\"Domain Labels:\", labels.tolist())\n        break","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-20T15:43:31.496766Z","iopub.execute_input":"2025-03-20T15:43:31.497037Z","iopub.status.idle":"2025-03-20T15:43:31.503347Z","shell.execute_reply.started":"2025-03-20T15:43:31.497007Z","shell.execute_reply":"2025-03-20T15:43:31.502698Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile Gen_turbo.py\n\nimport torch\nimport torch.nn as nn\nfrom torchinfo import summary\n\n\nclass Gen(nn.Module):\n    def __init__(self,image_channel,feature_g=64,noise_std=0.2):\n        super().__init__()\n        self.noise_std = noise_std\n        self.down1 = nn.Sequential(\n            nn.Conv2d(image_channel+1,feature_g,4,2,1),\n            nn.LeakyReLU(0.2,inplace=True)\n        )\n        self.down2 = self.down_block_conv(feature_g,feature_g*2,4,2,1)\n        self.down3 = self.down_block_conv(feature_g*2, feature_g*4, 4, 2, 1)\n        self.down4 = self.down_block_conv(feature_g*4, feature_g*8, 4, 2, 1)\n        self.down5 = self.down_block_conv(feature_g*8, feature_g*8, 4, 2, 1)\n        self.down6 = self.down_block_conv(feature_g*8, feature_g*8, 4, 2, 1)\n        self.down7 = self.down_block_conv(feature_g*8, feature_g*8, 4, 2, 1)\n\n        self.bottom = nn.Sequential(\n            nn.Conv2d(feature_g*8, feature_g*8, 4, 2, 1),\n            nn.LeakyReLU(0.2,inplace=True)\n        )\n\n        self.up1 = self.up_block_conv(feature_g*8,feature_g*8,4,  2,  1,drop_out=True)\n        self.up2 = self.up_block_conv(feature_g*8*2,feature_g*8,4, 2, 1,drop_out=True)\n        self.up3 = self.up_block_conv(feature_g*8*2,feature_g*8,4, 2, 1,drop_out=True)\n        self.up4 = self.up_block_conv(feature_g*8*2,feature_g*8,4, 2, 1)\n        self.up5 = self.up_block_conv(feature_g*8*2,feature_g*4,4, 2, 1)\n        self.up6 = self.up_block_conv(feature_g*4*2,feature_g*2,4, 2, 1)\n        self.up7 = self.up_block_conv(feature_g*2*2,feature_g,4, 2, 1)\n\n        self.final = nn.Sequential(\n            nn.ConvTranspose2d(feature_g,3,4,2,1),\n            nn.Tanh()\n        )\n\n    def up_block_conv(self,in_channel,out_channel,kernel,stride,padding,drop_out=False):\n        return nn.Sequential(\n            nn.ConvTranspose2d(in_channel,out_channel,kernel,stride,padding,bias=False),\n            nn.InstanceNorm2d(out_channel, affine=True),\n            nn.Dropout(0.5) if drop_out else nn.Identity(),\n            nn.ReLU(inplace=True)\n        )\n\n    def down_block_conv(self,in_channel,out_channel,kernel,stride,padding):\n        return nn.Sequential(\n            nn.Conv2d(in_channel, out_channel, kernel, stride, padding,padding_mode='reflect',bias=False),\n            nn.InstanceNorm2d(out_channel, affine=True),\n            nn.LeakyReLU(0.2,inplace=True)\n        )\n\n    def forward(self, x, direction):\n        condition = torch.ones_like(x[:, :1, :, :]) * direction.view(-1, 1, 1, 1)\n        x = torch.cat([x, condition], dim=1)\n        d1 = self.down1(x)\n        d2 = self.down2(d1)\n        d3 = self.down3(d2)\n        d4 = self.down4(d3)\n        d5 = self.down5(d4)\n        d6 = self.down6(d5)\n        d7 = self.down7(d6)\n\n        bottom = self.bottom(d7)\n\n        u1 = self.up1(bottom)\n        u2 = self.up2(torch.cat([u1, d7], 1))\n        u3 = self.up3(torch.cat([u2, d6], 1))\n        u4 = self.up4(torch.cat([u3, d5], 1))\n        u5 = self.up5(torch.cat([u4, d4], 1))\n        u6 = self.up6(torch.cat([u5, d3], 1))\n        u7 = self.up7(torch.cat([u6, d2], 1))\n        final = self.final(u7)\n\n        return final\n\nif __name__ == '__main__':\n    model = Gen(3,64)\n    print(summary(model,input_size=(1,3,256,256)))\n\n    # model.eval()\n    # dummy_input = (torch.randn(1, 3, 256, 256, dtype=torch.float32),)\n    # torch.onnx.export(model, dummy_input, \"model.onnx\", opset_version=11)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-20T15:43:31.504036Z","iopub.execute_input":"2025-03-20T15:43:31.504217Z","iopub.status.idle":"2025-03-20T15:43:31.523022Z","shell.execute_reply.started":"2025-03-20T15:43:31.504201Z","shell.execute_reply":"2025-03-20T15:43:31.522257Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile Dis_turbo.py\n\nimport torch\nimport torch.nn as nn\nfrom torchinfo import summary\n\nclass GaussianNoise(nn.Module):\n    def __init__(self, std=0.2):\n        super(GaussianNoise, self).__init__()\n        self.std = std\n\n    def forward(self, x):\n        if self.training:\n            noise = torch.randn_like(x) * self.std\n            return x + noise\n        return x\n\nclass Dis(nn.Module):\n    def __init__(self, image_channel=3, feature_d=64, num_classes=10, noise_std=0.2):\n        super().__init__()\n        self.noise_std = noise_std\n        self.conv1 = nn.Sequential(\n            nn.Conv2d(image_channel, feature_d, 4, 2, 1),\n            nn.LeakyReLU(0.2, inplace=True)\n        )\n        self.conv2 = self.block_conv(feature_d, feature_d * 2, 4, 2, 1)\n        self.conv3 = self.block_conv(feature_d * 2, feature_d * 4, 4, 2, 1)\n        self.conv4 = self.block_conv(feature_d * 4, feature_d * 8, 4, 1, 1)\n\n        self.dis_out = nn.Conv2d(feature_d * 8, 1, 4, 1, 1)\n        self.classifier = nn.Sequential(\n            nn.Conv2d(feature_d * 8, num_classes, 4, 1, 1),\n            nn.Flatten(),\n        )\n\n    def block_conv(self, in_channel, out_channel, kernel, stride, padding):\n        return nn.Sequential(\n            nn.Conv2d(in_channel, out_channel, kernel_size=kernel, stride=stride, padding=padding, bias=False),\n            nn.InstanceNorm2d(out_channel, affine=True),\n            nn.LeakyReLU(0.2, inplace=True)\n        )\n\n    def forward(self, x, return_feature=False):\n        x = self.conv1(x)\n        x = self.conv2(x)\n        x = self.conv3(x)\n        f_x = self.conv4(x)\n\n        dis_out = self.dis_out(f_x)\n        cls_out = self.classifier(f_x)\n\n        if return_feature:\n            return dis_out, cls_out, f_x\n        return dis_out, cls_out\n\nif __name__ == '__main__':\n    model = Dis(image_channel=3, feature_d=64, num_classes=10)\n    print(summary(model, input_size=(32, 3, 256, 256)))\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-20T15:43:31.685598Z","iopub.execute_input":"2025-03-20T15:43:31.686018Z","iopub.status.idle":"2025-03-20T15:43:31.69106Z","shell.execute_reply.started":"2025-03-20T15:43:31.685981Z","shell.execute_reply":"2025-03-20T15:43:31.690295Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile loss_turbo.py\n\nimport torch\nimport torch.nn as nn\n\ndef cycle_loss_with_feature(dis_x, real_x, reconstructed_x, gamma, l1 = nn.L1Loss()):\n    _,_,feature_real = dis_x(real_x.detach(), return_feature=True)\n    _,_,feature_fake = dis_x(reconstructed_x.detach(), return_feature=True)\n    feature_loss = l1(feature_fake,feature_real)\n    cycle_loss = l1(reconstructed_x,real_x)\n    loss = gamma * feature_loss + (1 - gamma) * cycle_loss\n    return loss\n\ndef compute_discriminator_loss(discriminator, real_image, fake_image, labels, smooth_label=0.0,\n                               bce = nn.BCEWithLogitsLoss(),\n                               mce = nn.CrossEntropyLoss()):\n\n    out_d_real, out_d_classifier_real = discriminator(real_image)\n    loss_d_real = bce(out_d_real, torch.ones_like(out_d_real) * (1-smooth_label))\n    out_d_fake, out_d_classifier_fake = discriminator(fake_image.detach())\n    loss_d_fake = bce(out_d_fake, torch.zeros_like(out_d_fake) + smooth_label)\n    loss_class_real = mce(out_d_classifier_real, labels)\n    loss_class_fake = mce(out_d_classifier_fake, labels)\n\n    loss_d = (loss_d_real + loss_d_fake + loss_class_real + loss_class_fake) / 2\n    return loss_d, loss_d_real, loss_d_fake, loss_class_real, loss_class_fake\n\ndef compute_generator_loss(discriminator,\n                           true_image, gen_image, cycle_image, identity_image,\n                           labels,\n                           cycle_lambda, identity_lambda, classifier_lambda,\n                           gamma,\n                           bce = nn.BCEWithLogitsLoss(),\n                           mce = nn.CrossEntropyLoss(),\n                           l1 = nn.L1Loss()):\n\n    out_d, out_d_classifier = discriminator(gen_image)\n    loss_g = bce(out_d, torch.ones_like(out_d))\n    loss_class = mce(out_d_classifier, labels)\n    loss_iden = l1(identity_image, true_image)\n    loss_cycle = cycle_lambda * cycle_loss_with_feature(discriminator, true_image, cycle_image, gamma=gamma)\n\n    loss_g_final = loss_g + loss_cycle + identity_lambda * loss_iden + classifier_lambda * loss_class\n    return loss_g_final","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-20T15:43:31.691858Z","iopub.execute_input":"2025-03-20T15:43:31.692118Z","iopub.status.idle":"2025-03-20T15:43:31.708805Z","shell.execute_reply.started":"2025-03-20T15:43:31.692087Z","shell.execute_reply":"2025-03-20T15:43:31.708023Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile callback.py\n\nfrom torch.optim.lr_scheduler import LambdaLR\n\n\ndef get_lambda_scheduler(epoch, total_epochs, decay_start_epoch, start_lambda, end_lambda, mode='linear'):\n    if epoch < decay_start_epoch:\n        return start_lambda\n\n    if mode == 'linear':\n        decay_epochs = total_epochs - decay_start_epoch\n        return max(start_lambda - (start_lambda - end_lambda) * ((epoch - decay_start_epoch) / float(decay_epochs)),\n                   end_lambda)\n\n    if mode == 'linear_reverse':\n        decay_epochs = total_epochs - decay_start_epoch\n        return min(start_lambda + (end_lambda - start_lambda) * ((epoch - decay_start_epoch) / float(decay_epochs)),\n                   end_lambda)\n\n    elif mode == 'exp':\n        decay_rate = (end_lambda / start_lambda) ** (1 / (total_epochs - decay_start_epoch))\n        return max(start_lambda * (decay_rate ** (epoch - decay_start_epoch)), end_lambda)\n    else:\n        raise ValueError(\"Mode not found !!!\")\n\ndef get_lr_scheduler(optimizer, total_epochs, decay_start_epoch):\n    def lambda_rule(epoch):\n        return 1.0 - max(0, epoch - decay_start_epoch) / float(total_epochs - decay_start_epoch)\n    return LambdaLR(optimizer, lr_lambda=lambda_rule)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-20T15:43:31.710183Z","iopub.execute_input":"2025-03-20T15:43:31.710455Z","iopub.status.idle":"2025-03-20T15:43:31.725953Z","shell.execute_reply.started":"2025-03-20T15:43:31.710428Z","shell.execute_reply":"2025-03-20T15:43:31.725325Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nos.environ[\"TF_ENABLE_ONEDNN_OPTS\"] = \"0\"\nimport shutil\nimport torch\nimport torchvision.utils\nfrom Dis_turbo import Dis\nfrom Gen_turbo import Gen\nfrom Gen_turbo import Gen\nfrom torch.utils.data import DataLoader, Subset\nfrom dataset_turbo import MultiPathStyleDataset\nfrom torchvision.transforms import ToTensor, Resize, Normalize, RandomCrop\nfrom torchvision.transforms.v2 import Compose\nfrom tqdm import tqdm\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.tensorboard import SummaryWriter\nfrom argparse import ArgumentParser\nimport random\nfrom callback import get_lambda_scheduler,get_lr_scheduler\nfrom loss_turbo import cycle_loss_with_feature,compute_discriminator_loss,compute_generator_loss\n\n\ndef get_arg():\n    parser = ArgumentParser(description='Train Cycle GAN turbo')\n    parser.add_argument('--epochs', '-e', type=int, default=100)\n    parser.add_argument('--batch_size', '-b', type=int, default=4)\n    parser.add_argument('--learning_rate', '-lr', type=float, default=0.0002)\n    parser.add_argument('--ratio_train_size', '-rt', type=float, default=1)\n    parser.add_argument('--val_size', '-r', type=int, default=32)\n    parser.add_argument('--cycle_lambda', '-cyl', type=int, default=5)\n    parser.add_argument('--identity_lambda', '-il', type=int, default=5)\n    parser.add_argument('--classifier_lambda', '-pl', type=int, default=1)\n    parser.add_argument('--checkpoint', '-c', type=str, default=None)\n    parser.add_argument('--trained', '-t', type=str, default='trained_image2paint')\n    parser.add_argument('--log', '-l', type=str, default='tensorboard')\n    args, unknown = parser.parse_known_args()\n    return args\n\ndef split_data(data, ratio_train_size, val_size, seed=42):\n    random.seed(seed)\n    index_data = list(range(len(data)))\n    random.shuffle(index_data)\n    train_size = int(ratio_train_size * len(index_data))\n    train_index = index_data[:train_size]\n    val_index = index_data[train_size:train_size + val_size]\n\n    train = Subset(data, train_index)\n    val = Subset(data, val_index)\n    print(f'Training set size: {len(train)}')\n    print(f'Validation set size: {len(val)}')\n    return train, val\n\ndef train(arg):\n    if not os.path.isdir(arg.trained):\n        os.mkdir(arg.trained)\n    if os.path.isdir(arg.log):\n        shutil.rmtree(arg.log)\n\n    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n    writer = SummaryWriter(arg.log)\n\n    raw_transform = Compose([\n        Resize(286),\n        RandomCrop(256),\n        ToTensor(),\n        Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])\n    ])\n\n    style_transform = Compose([\n        Resize(286),\n        RandomCrop(256),\n        ToTensor(),\n        Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])\n    ])\n\n    style_folder = [\n        '/kaggle/input/monet2photo/trainA',\n        '/kaggle/input/vangogh2photo/vangogh2photo/trainA',\n        #\"D:/ukiyoe_painting\",\n        #\"D:/cezanne_painting\"\n    ]\n    data = MultiPathStyleDataset(\n        raw_root=\"/kaggle/input/gan-getting-started/photo_jpg\",\n        style_folders=style_folder,\n        raw_transform=raw_transform,\n        style_transform=style_transform)\n\n    train,val = split_data(data,arg.ratio_train_size,arg.val_size)\n    train_loader = DataLoader(train, batch_size=arg.batch_size, shuffle=True,num_workers=4)\n    val_loader = DataLoader(val, batch_size=arg.val_size, shuffle=False,num_workers=4)\n\n    gen = Gen(3,64)\n    dis = Dis(3,64)\n\n    gen.to(device)\n    dis.to(device)\n\n    def init_weights(m,mode='standard'):\n        if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):\n            if mode=='standard':\n                nn.init.normal_(m.weight.data, 0.0, 0.02)\n            else:\n                nn.init.kaiming_normal_(m.weight)\n\n    gen.apply(init_weights)\n    dis.apply(init_weights)\n\n    op_g = optim.Adam(gen.parameters(), lr=arg.learning_rate, betas=(0.5, 0.999))\n    op_d = optim.Adam(dis.parameters(), lr=arg.learning_rate, betas=(0.5, 0.999))\n\n    scheduler_g = get_lr_scheduler(op_g, total_epochs=arg.epochs, decay_start_epoch=arg.epochs / 2)\n    scheduler_d = get_lr_scheduler(op_d, total_epochs=arg.epochs, decay_start_epoch=arg.epochs / 2)\n\n    if arg.checkpoint:\n        checkpoint = torch.load(arg.checkpoint,map_location=lambda storage, loc: storage.cuda(torch.cuda.current_device()),weights_only=True)\n        start_epoch = checkpoint['epoch']\n        dis.load_state_dict(checkpoint['dis'])\n        gen.load_state_dict(checkpoint['gen'])\n        op_g.load_state_dict(checkpoint['op_g'])\n        op_d.load_state_dict(checkpoint['op_d'])\n        best_loss = checkpoint['best_loss']\n    else:\n        start_epoch = 0\n        best_loss = float('inf')\n\n    train_iters = len(train) // arg.batch_size\n    val_iters = len(val) // arg.batch_size\n\n    for epoch in range(start_epoch, arg.epochs):\n        dis.train()\n        gen.train()\n\n        mean_loss_d = 0.0\n        mean_loss_d_ori = 0.0\n        mean_loss_d_class = 0.0\n        mean_loss_ga = 0.0\n        mean_loss_gb = 0.0\n\n        mean_loss_ga_perceptual = 0.0\n        mean_loss_gb_perceptual = 0.0\n\n        cycle_lambda = get_lambda_scheduler(epoch,arg.epochs,0,\n                                            arg.cycle_lambda,arg.cycle_lambda/arg.cycle_lambda,\n                                            mode='linear')\n        gamma = get_lambda_scheduler(epoch,arg.epochs,0,\n                                     0.0001,0.999,\n                                     mode='linear_reverse')\n\n        for raw_image, true_image, labels in tqdm(train_loader):\n            raw_image = raw_image.to(device)\n            true_image = true_image.to(device)\n            labels = labels.to(device)\n\n            out_a = gen(raw_image, labels)\n            out_b = gen(true_image, torch.zeros_like(labels))\n            cycle_a = gen(out_a.detach(), labels)\n            cycle_b = gen(out_b.detach(), torch.zeros_like(labels))\n            iden_a = gen(true_image, labels)\n            iden_b = gen(raw_image,torch.zeros_like(labels))\n\n            (loss_da,\n             loss_da_real,loss_da_fake,\n             loss_class_real_a,loss_class_fake_a) = (compute_discriminator_loss(dis,\n                                                                        true_image,out_a,labels,\n                                                                        0.05))\n\n            (loss_db,\n             loss_db_real, loss_db_fake,\n             loss_class_real_b, loss_class_fake_b) = (compute_discriminator_loss(dis,\n                                                                         raw_image,out_b,torch.zeros_like(labels),\n                                                                         0))\n\n            mean_loss_d += loss_da.item() + loss_db.item()\n            mean_loss_d_ori += loss_da_real.item() + loss_da_fake.item() + loss_db_real.item() + loss_db_fake.item()\n            mean_loss_d_class += loss_class_real_a.item() + loss_class_fake_a.item() + loss_class_real_b.item() + loss_class_fake_b.item()\n\n            op_d.zero_grad()\n            loss_d = loss_da + loss_db\n            loss_d.backward()\n            op_d.step()\n\n            loss_ga_final = compute_generator_loss(dis,\n                                                   true_image,out_a,cycle_a,iden_a,labels,\n                                                   cycle_lambda,arg.identity_lambda,arg.classifier_lambda,\n                                                   gamma)\n\n            loss_gb_final = compute_generator_loss(dis,\n                                                   raw_image,out_b,cycle_b,iden_b,torch.zeros_like(labels),\n                                                   cycle_lambda,arg.identity_lambda,arg.classifier_lambda,\n                                                   gamma)\n\n            op_g.zero_grad()\n            loss_g_final = loss_ga_final + loss_gb_final\n            loss_g_final.backward()\n            op_g.step()\n\n        scheduler_g.step()\n        scheduler_d.step()\n\n        print(f'Epoch {epoch + 1}/{arg.epochs} | '\n              f'D - Loss: {mean_loss_d / train_iters:.4f} | '\n              f'ori loss: {mean_loss_d_ori / train_iters:.4f} - classifier loss: {mean_loss_d_class/ train_iters:.4f} | '\n              f'Ga - Loss: {mean_loss_ga / train_iters:.4f} | '\n              f'Ga - P Loss: {mean_loss_ga_perceptual / train_iters:.4f} | '\n              f'Gb - Loss: {mean_loss_gb / train_iters:.4f} | '\n              f'Gb - P Loss: {mean_loss_gb_perceptual / train_iters:.4f}')\n        print('Cycle: {:.2f} - Gamma: {:.2f}'.format(cycle_lambda,gamma))\n\n        gen.eval()\n        for raw_image, true_image, labels in val_loader:\n            raw_image = raw_image.to(device)\n            true_image = true_image.to(device)\n            labels = labels.to(device)\n            with torch.no_grad():\n                out_ga = gen(raw_image,labels)\n                out_gb = gen(out_ga,torch.zeros_like(labels))\n                true_image = torchvision.utils.make_grid(true_image[:arg.val_size].cpu(), normalize=True)\n                gen_image = torchvision.utils.make_grid(out_ga[:arg.val_size].cpu(), normalize=True)\n                raw_image = torchvision.utils.make_grid(raw_image[:arg.val_size].cpu(), normalize=True)\n                cycle_image = torchvision.utils.make_grid(out_gb[:arg.val_size].cpu(), normalize=True)\n\n                writer.add_image('true image', true_image, epoch)\n                writer.add_image('gen image', gen_image, epoch)\n                writer.add_image('raw image', raw_image, epoch)\n                writer.add_image('cycle image', cycle_image, epoch)\n\n                writer.add_scalar('loss d', mean_loss_d / train_iters, epoch)\n                writer.add_scalar('loss ga', mean_loss_ga / train_iters, epoch)\n                writer.add_scalar('loss gb', mean_loss_gb / train_iters, epoch)\n\n        checkpoint = {\n                'epoch': epoch + 1,\n                'gen': gen.state_dict(),\n                'dis': dis.state_dict(),\n                'op_g': op_g.state_dict(),\n                'op_d': op_d.state_dict(),\n                'best_loss': best_loss\n            }\n        torch.save(checkpoint, '{}/last.pt'.format(arg.trained))\n        if mean_loss_ga/train_iters < best_loss:\n            best_loss = mean_loss_ga/train_iters\n            torch.save(checkpoint, '{}/best.pt'.format(arg.trained))\n\nif __name__ == '__main__':\n    arg = get_arg()\n    train(arg)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-20T15:44:34.252985Z","iopub.execute_input":"2025-03-20T15:44:34.253279Z","iopub.status.idle":"2025-03-20T15:44:41.144567Z","shell.execute_reply.started":"2025-03-20T15:44:34.253256Z","shell.execute_reply":"2025-03-20T15:44:41.143256Z"}},"outputs":[],"execution_count":null}]}