{"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":"none","dataSources":[{"sourceId":21755,"databundleVersionId":1475600,"sourceType":"competition"},{"sourceId":1217826,"sourceType":"datasetVersion","datasetId":298806},{"sourceId":3956508,"sourceType":"datasetVersion","datasetId":850761}],"dockerImageVersionId":30786,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import torch\nimport torch.nn as nn\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","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **Configuration**","metadata":{}},{"cell_type":"code","source":"def save_checkpoint(model, optimizer, filename=\"/kaggle/working/.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\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    # If we don't do this then it will just have learning rate of old checkpoint\n    # and it will lead to many hours of debugging \\:\n    for param_group in optimizer.param_groups:\n        param_group[\"lr\"] = lr\n\n\ndef seed_everything(seed=42):\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\n\nDEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n#TRAIN_DIR = \"data/train\"\n#VAL_DIR = \"data/val\"\nBATCH_SIZE = 1\nLEARNING_RATE = 1e-5\nLAMBDA_IDENTITY = 0.0\nLAMBDA_CYCLE = 10\nNUM_WORKERS = 4\nNUM_EPOCHS = 100\nLOAD_MODEL = False\nSAVE_MODEL = True\nCHECKPOINT_GEN_A = \"/kaggle/working/genA.pth.tar\"\nCHECKPOINT_GEN_M = \"/kaggle/working/genM.pth.tar\"\nCHECKPOINT_CRITIC_A = \"/kaggle/working/criticA.pth.tar\"\nCHECKPOINT_CRITIC_M = \"/kaggle/working/criticM.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":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **Generator**","metadata":{}},{"cell_type":"code","source":"class ConvBlock(nn.Module):\n    def __init__(self, in_channels, out_channels, down=True, use_act=True, **kwargs):\n        super().__init__()\n        self.conv = nn.Sequential(\n            nn.Conv2d(in_channels, out_channels, padding_mode=\"reflect\", **kwargs)\n            if down\n            else nn.ConvTranspose2d(in_channels, out_channels, **kwargs),\n            nn.InstanceNorm2d(out_channels),\n            nn.ReLU(inplace=True) if use_act else nn.Identity(),\n        )\n\n    def forward(self, x):\n        return self.conv(x)\n\n\nclass ResidualBlock(nn.Module):\n    def __init__(self, channels):\n        super().__init__()\n        self.block = nn.Sequential(\n            ConvBlock(channels, channels, kernel_size=3, padding=1),\n            ConvBlock(channels, channels, use_act=False, kernel_size=3, padding=1),\n        )\n\n    def forward(self, x):\n        return x + self.block(x)\n\n\nclass Generator(nn.Module):\n    def __init__(self, img_channels, num_features=64, num_residuals=9):\n        super().__init__()\n        self.initial = nn.Sequential(\n            nn.Conv2d(\n                img_channels,\n                num_features,\n                kernel_size=7,\n                stride=1,\n                padding=3,\n                padding_mode=\"reflect\",\n            ),\n            nn.InstanceNorm2d(num_features),\n            nn.ReLU(inplace=True),\n        )\n        self.down_blocks = nn.ModuleList(\n            [\n                ConvBlock(\n                    num_features, num_features * 2, kernel_size=3, stride=2, padding=1\n                ),\n                ConvBlock(\n                    num_features * 2,\n                    num_features * 4,\n                    kernel_size=3,\n                    stride=2,\n                    padding=1,\n                ),\n            ]\n        )\n        self.res_blocks = nn.Sequential(\n            *[ResidualBlock(num_features * 4) for _ in range(num_residuals)]\n        )\n        self.up_blocks = nn.ModuleList(\n            [\n                ConvBlock(\n                    num_features * 4,\n                    num_features * 2,\n                    down=False,\n                    kernel_size=3,\n                    stride=2,\n                    padding=1,\n                    output_padding=1,\n                ),\n                ConvBlock(\n                    num_features * 2,\n                    num_features * 1,\n                    down=False,\n                    kernel_size=3,\n                    stride=2,\n                    padding=1,\n                    output_padding=1,\n                ),\n            ]\n        )\n\n        self.last = nn.Conv2d(\n            num_features * 1,\n            img_channels,\n            kernel_size=7,\n            stride=1,\n            padding=3,\n            padding_mode=\"reflect\",\n        )\n\n    def forward(self, x):\n        x = self.initial(x)\n        for layer in self.down_blocks:\n            x = layer(x)\n        x = self.res_blocks(x)\n        for layer in self.up_blocks:\n            x = layer(x)\n        return torch.tanh(self.last(x))\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **Discriminator**","metadata":{}},{"cell_type":"code","source":"class Block(nn.Module):\n    def __init__(self, in_channels, out_channels, stride):\n        super().__init__()\n        self.conv = nn.Sequential(\n            nn.Conv2d(\n                in_channels,\n                out_channels,\n                4,\n                stride,\n                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)\n\n\nclass Discriminator(nn.Module):\n    def __init__(self, in_channels=3, features=[64, 128, 256, 512]):\n        super().__init__()\n        self.initial = nn.Sequential(\n            nn.Conv2d(\n                in_channels,\n                features[0],\n                kernel_size=4,\n                stride=2,\n                padding=1,\n                padding_mode=\"reflect\",\n            ),\n            nn.LeakyReLU(0.2, inplace=True),\n        )\n\n        layers = []\n        in_channels = features[0]\n        for feature in features[1:]:\n            layers.append(\n                Block(in_channels, feature, stride=1 if feature == features[-1] else 2)\n            )\n            in_channels = feature\n        layers.append(\n            nn.Conv2d(\n                in_channels,\n                1,\n                kernel_size=4,\n                stride=1,\n                padding=1,\n                padding_mode=\"reflect\",\n            )\n        )\n        self.model = nn.Sequential(*layers)\n\n    def forward(self, x):\n        x = self.initial(x)\n        return torch.sigmoid(self.model(x))\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **Buliding**","metadata":{}},{"cell_type":"code","source":"class MonetToActual(Dataset):\n    def __init__(self, root_monet, root_actual, transform=None):\n        self.root_monet = root_monet\n        self.root_actual = root_actual\n        self.transform = transform\n\n        self.monet_images = os.listdir(root_monet)\n        self.actual_images = os.listdir(root_actual)[:1000]\n        self.length_dataset = max(len(self.monet_images), len(self.actual_images)) # 1000, 1500\n        self.actual_len = len(self.actual_images)\n        self.monet_len = len(self.monet_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        actual_img = self.actual_images[index % self.actual_len]\n\n        actual_path = os.path.join(self.root_actual, actual_img)\n        monet_path = os.path.join(self.root_monet, monet_img)\n\n        monet_img = np.array(Image.open(monet_path).convert(\"RGB\"))\n        actual_img = np.array(Image.open(actual_path).convert(\"RGB\"))\n\n        if self.transform:\n            augmentations = self.transform(image=monet_img, image0=actual_img)\n            monet_img = augmentations[\"image\"]\n            actual_img = augmentations[\"image0\"]\n\n        return monet_img, actual_img","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **Loading Data**","metadata":{}},{"cell_type":"code","source":"dataset = MonetToActual(\n        root_monet='/kaggle/input/gan-getting-started/monet_jpg',\n        root_actual='/kaggle/input/gan-getting-started/photo_jpg',\n        transform=transforms,\n    )\nloader = DataLoader(\n        dataset,\n        batch_size=BATCH_SIZE,\n        shuffle=True,\n        num_workers=NUM_WORKERS,\n        pin_memory=True,\n    )","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **Training Function**","metadata":{}},{"cell_type":"code","source":"def train_fn(\n    disc_A, disc_M, gen_M, gen_A, loader, opt_disc, opt_gen, l1, mse, d_scaler, g_scaler,epoch\n):\n    H_reals = 0\n    H_fakes = 0\n    A_reals=0\n    A_fakes=0\n    loop = tqdm(loader, leave=True)\n\n    for idx, (monet, actual) in enumerate(loop):\n        monet = monet.to(DEVICE)\n        actual = actual.to(DEVICE)\n\n        # Train Discriminators A and M\n        with torch.cuda.amp.autocast():\n            fake_actual = gen_A(monet)\n            D_A_real = disc_A(actual)\n            D_A_fake = disc_A(fake_actual.detach())\n            A_reals += D_A_real.mean().item()\n            A_fakes += D_A_fake.mean().item()\n            D_A_real_loss = mse(D_A_real, torch.ones_like(D_A_real))\n            D_A_fake_loss = mse(D_A_fake, torch.zeros_like(D_A_fake))\n            D_A_loss = D_A_real_loss + D_A_fake_loss\n\n            fake_monet = gen_M(actual)\n            D_M_real = disc_M(monet)\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            # put it togethor\n            D_loss = (D_A_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#H->A\n#Z->M\n        # Train Generators H and Z\n        with torch.cuda.amp.autocast():\n            # adversarial loss for both generators\n            D_A_fake = disc_A(fake_actual)\n            D_M_fake = disc_M(fake_monet)\n            loss_G_A = mse(D_A_fake, torch.ones_like(D_A_fake))\n            loss_G_M = mse(D_M_fake, torch.ones_like(D_M_fake))\n\n            # cycle loss\n            cycle_monet = gen_M(fake_actual)\n            cycle_actual = gen_A(fake_monet)\n            cycle_monet_loss = l1(monet, cycle_monet)\n            cycle_actual_loss = l1(actual, cycle_actual)\n\n            # identity loss (remove these for efficiency if you set lambda_identity=0)\n            identity_monet = gen_M(monet)\n            identity_actual = gen_A(actual)\n            identity_monet_loss = l1(monet, identity_monet)\n            identity_actual_loss = l1(actual, identity_actual)\n\n            # add all togethor\n            G_loss = (\n                loss_G_M\n                + loss_G_A\n                + cycle_monet_loss * LAMBDA_CYCLE\n                + cycle_actual_loss * LAMBDA_CYCLE\n                + identity_actual_loss * LAMBDA_IDENTITY\n                + identity_monet_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        if idx % 200 == 0:\n            save_image(fake_actual * 0.5 + 0.5, f\"/kaggle/working/actual_{idx}_{epoch}.png\")\n            save_image(fake_monet * 0.5 + 0.5, f\"/kaggle/working/monet_{idx}_{epoch}.png\")\n\n        loop.set_postfix(A_real=A_reals / (idx + 1), A_fake=A_fakes / (idx + 1))","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def main():\n    disc_A = 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_A = Generator(img_channels=3, num_residuals=9).to(DEVICE)\n    opt_disc = optim.Adam(\n        list(disc_A.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_A.parameters()),\n        lr=LEARNING_RATE,\n        betas=(0.5, 0.999),\n    )\n\n    L1 = nn.L1Loss()\n    mse = nn.MSELoss()\n\n\n    dataset = MonetToActual(\n        root_monet='/kaggle/input/gan-getting-started/monet_jpg',\n        root_actual='/kaggle/input/gan-getting-started/photo_jpg',\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    for epoch in range(NUM_EPOCHS):\n        train_fn(\n            disc_A,\n            disc_M,\n            gen_M,\n            gen_A,\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_A, opt_gen, filename=CHECKPOINT_GEN_A)\n        save_checkpoint(gen_M, opt_gen, filename=CHECKPOINT_GEN_M)\n        save_checkpoint(disc_A, opt_disc, filename=CHECKPOINT_CRITIC_A)\n        save_checkpoint(disc_M, opt_disc, filename=CHECKPOINT_CRITIC_M)\n\nmain()","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}