{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# **CycleGAN**","metadata":{}},{"cell_type":"code","source":"!pip install deepspeed","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import glob\nimport os\nimport shutil\nimport deepspeed as ds\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport pytorch_lightning as L\nimport torch\nimport torch.nn.functional as F\nimport torchvision.transforms as T\nfrom pytorch_lightning.utilities import CombinedLoader\nfrom torch import nn\nfrom torch.utils.data import DataLoader, Dataset\nfrom torchvision.io import read_image\nfrom torchvision.utils import make_grid, save_image\n_ = L.seed_everything(0, workers=True)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def show_img(img_tensor, nrow, title=\"\"):\n    img_tensor = img_tensor.detach().cpu() * 0.5 + 0.5\n    img_grid = make_grid(img_tensor, nrow=nrow).permute(1, 2, 0)\n    plt.figure(figsize=(10, 7))\n    plt.imshow(img_grid)\n    plt.axis(\"off\")\n    plt.title(title)\n    plt.show()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CustomTransform(object):\n    def __init__(self, load_dim=286, target_dim=256):\n        self.transform_train = T.Compose([\n            T.Resize((load_dim, load_dim), antialias=True),\n            T.RandomCrop((target_dim, target_dim)),\n            T.RandomHorizontalFlip(p=0.5),\n            T.ColorJitter(brightness=0.2, contrast=0.2,\n                          saturation=0.2, hue=0.1),\n        ])\n        self.transform = T.Resize((target_dim, target_dim), antialias=True)   \n    def __call__(self, img, stage):\n        if stage == \"fit\":\n            img = self.transform_train(img)\n        else:\n            img = self.transform(img)\n        return img * 2 - 1","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CustomDataset(Dataset):\n    def __init__(self, filenames, transform, stage):\n        self.filenames = filenames\n        self.transform = transform\n        self.stage = stage\n        \n    def __len__(self):\n        return len(self.filenames)\n    \n    def __getitem__(self, idx):\n        img_name = self.filenames[idx]\n        img = read_image(img_name) / 255.0\n        return self.transform(img, stage=self.stage)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"DEBUG = False\n\nDM_CONFIG = {    \n    \"monet_dir\": os.path.join(\"/kaggle/input/gan-getting-started/monet_jpg\", \"*.jpg\"),\n    \"photo_dir\": os.path.join(\"/kaggle/input/gan-getting-started/photo_jpg\", \"*.jpg\"),\n    \n    \"loader_config\": {\n        \"num_workers\": os.cpu_count(),\n        \"pin_memory\": torch.cuda.is_available(),\n    },\n    \"sample_size\": 5,\n    \"batch_size\": 1 if not DEBUG else 1,\n}","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CustomDataModule(L.LightningDataModule):\n    def __init__(\n        self,\n        monet_dir,\n        photo_dir, \n        loader_config,\n        sample_size,\n        batch_size,\n    ):\n        super().__init__()\n        self.loader_config = loader_config\n        self.sample_size = sample_size\n        self.batch_size = batch_size       \n            \n        # store file paths\n        self.monet_filenames = sorted(glob.glob(monet_dir))\n        self.photo_filenames = sorted(glob.glob(photo_dir))\n        \n        # define transformations for image augmentation\n        self.transform = CustomTransform()\n        \n    def setup(self, stage):\n        if stage == \"fit\":\n            self.train_monet = CustomDataset(self.monet_filenames, self.transform, stage)\n            self.train_photo = CustomDataset(self.photo_filenames, self.transform, stage)\n            \n        if stage in [\"fit\", \"test\", \"predict\"]:\n            # to be used for test and prediction datasets as well\n            self.valid_photo = CustomDataset(self.photo_filenames, self.transform, None)\n            \n    def train_dataloader(self):\n        loader_config = {\n            \"shuffle\": True,\n            \"drop_last\": True,\n            \"batch_size\": self.batch_size,\n            **self.loader_config,\n        }\n        loader_monet = DataLoader(self.train_monet, **loader_config)\n        loader_photo = DataLoader(self.train_photo, **loader_config)\n        loaders = {\"monet\": loader_monet, \"photo\": loader_photo}\n        return CombinedLoader(loaders, mode=\"max_size_cycle\")\n    \n    def val_dataloader(self):\n        return DataLoader(\n            self.valid_photo,\n            batch_size=self.sample_size,\n            **self.loader_config,\n        )\n    \n    def test_dataloader(self):\n        return self.val_dataloader()\n    \n    def predict_dataloader(self):\n        return DataLoader(\n            self.valid_photo,\n            batch_size=self.batch_size,\n            **self.loader_config,\n        )","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"dm_sample = CustomDataModule(batch_size=5, **{k: v for k, v in DM_CONFIG.items() if k != \"batch_size\"})\ndm_sample.setup(\"fit\")\ntrain_loader = dm_sample.train_dataloader()\nimgs = next(iter(train_loader))\nshow_img(imgs[\"monet\"], nrow=5, title=\"Augmented Monet Paintings\")\nshow_img(imgs[\"photo\"], nrow=5, title=\"Augmented Photos\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Downsampling(nn.Module):\n    def __init__(\n        self,\n        in_channels,\n        out_channels,\n        kernel_size=4,\n        stride=2,\n        padding=1,\n        norm=True,\n        lrelu=True,\n    ):\n        super().__init__()\n        self.block = nn.Sequential(\n            nn.Conv2d(in_channels, out_channels,\n                      kernel_size=kernel_size, stride=stride, padding=padding, bias=not norm),\n        )\n        if norm:\n            self.block.append(nn.InstanceNorm2d(out_channels, affine=True))\n        if lrelu is not None:\n            self.block.append(nn.LeakyReLU(0.2, True) if lrelu else nn.ReLU(True))\n        \n    def forward(self, x):\n        return self.block(x)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Upsampling(nn.Module):\n    def __init__(\n        self,\n        in_channels,\n        out_channels,\n        kernel_size=4,\n        stride=2,\n        padding=1,\n        output_padding=0,\n        dropout=False,\n    ):\n        super().__init__()\n        self.block = nn.Sequential(\n            nn.ConvTranspose2d(in_channels, out_channels,\n                               kernel_size=kernel_size, stride=stride, \n                               padding=padding, output_padding=output_padding, bias=False),\n            nn.InstanceNorm2d(out_channels, affine=True),\n        )\n        if dropout:\n            self.block.append(nn.Dropout(0.5))\n        self.block.append(nn.ReLU(True))\n        \n    def forward(self, x):\n        return self.block(x)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class ResBlock(nn.Module):\n    def __init__(self, in_channels, kernel_size=3, padding=1):\n        super().__init__()\n        self.block = nn.Sequential(\n            nn.ReflectionPad2d(padding),\n            Downsampling(in_channels, in_channels,\n                         kernel_size=kernel_size, stride=1, padding=0, lrelu=False),\n            nn.ReflectionPad2d(padding),\n            Downsampling(in_channels, in_channels,\n                         kernel_size=kernel_size, stride=1, padding=0, lrelu=None),\n        )\n        \n    def forward(self, x):\n        return x + self.block(x)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class UNetGenerator(nn.Module):\n    def __init__(self, hid_channels, in_channels, out_channels):\n        super().__init__()\n        self.downsampling_path = nn.Sequential(\n            Downsampling(in_channels, hid_channels, norm=False), # 64x128x128\n            Downsampling(hid_channels, hid_channels*2), # 128x64x64\n            Downsampling(hid_channels*2, hid_channels*4), # 256x32x32\n            Downsampling(hid_channels*4, hid_channels*8), # 512x16x16\n            Downsampling(hid_channels*8, hid_channels*8), # 512x8x8\n            Downsampling(hid_channels*8, hid_channels*8), # 512x4x4\n            Downsampling(hid_channels*8, hid_channels*8), # 512x2x2\n            Downsampling(hid_channels*8, hid_channels*8, norm=False), # 512x1x1, instance norm does not work on 1x1\n        )\n        self.upsampling_path = nn.Sequential(\n            Upsampling(hid_channels*8, hid_channels*8, dropout=True), # (512+512)x2x2\n            Upsampling(hid_channels*16, hid_channels*8, dropout=True), # (512+512)x4x4\n            Upsampling(hid_channels*16, hid_channels*8, dropout=True), # (512+512)x8x8\n            Upsampling(hid_channels*16, hid_channels*8), # (512+512)x16x16\n            Upsampling(hid_channels*16, hid_channels*4), # (256+256)x32x32\n            Upsampling(hid_channels*8, hid_channels*2), # (128+128)x64x64\n            Upsampling(hid_channels*4, hid_channels), # (64+64)x128x128\n        )\n        self.feature_block = nn.Sequential(\n            nn.ConvTranspose2d(hid_channels*2, out_channels,\n                               kernel_size=4, stride=2, padding=1), # 3x256x256\n            nn.Tanh(),\n        )\n        \n    def forward(self, x):\n        skips = []\n        for down in self.downsampling_path:\n            x = down(x)\n            skips.append(x)\n        skips = reversed(skips[:-1])\n\n        for up, skip in zip(self.upsampling_path, skips):\n            x = up(x)\n            x = torch.cat([x, skip], dim=1)\n        return self.feature_block(x)\n    \nclass ResNetGenerator(nn.Module):\n    def __init__(self, hid_channels, in_channels, out_channels, num_resblocks):\n        super().__init__()\n        self.model = nn.Sequential(\n            nn.ReflectionPad2d(3),\n            Downsampling(in_channels, hid_channels,\n                         kernel_size=7, stride=1, padding=0, lrelu=False), # 64x256x256\n            Downsampling(hid_channels, hid_channels*2, kernel_size=3, lrelu=False), # 128x128x128\n            Downsampling(hid_channels*2, hid_channels*4, kernel_size=3, lrelu=False), # 256x64x64\n            *[ResBlock(hid_channels*4) for _ in range(num_resblocks)], # 256x64x64\n            Upsampling(hid_channels*4, hid_channels*2, kernel_size=3, output_padding=1), # 128x128x128\n            Upsampling(hid_channels*2, hid_channels, kernel_size=3, output_padding=1), # 64x256x256\n            nn.ReflectionPad2d(3),\n            nn.Conv2d(hid_channels, out_channels, kernel_size=7, stride=1, padding=0), # 3x256x256\n            nn.Tanh(),\n        )\n        \n    def forward(self, x):\n        return self.model(x)\n    \ndef get_gen(gen_name, hid_channels, num_resblocks, in_channels=3, out_channels=3):\n    if gen_name == \"unet\":\n        return UNetGenerator(hid_channels, in_channels, out_channels)\n    elif gen_name == \"resnet\":\n        return ResNetGenerator(hid_channels, in_channels, out_channels, num_resblocks)\n    else:\n        raise NotImplementedError(f\"Generator name '{gen_name}' not recognized.\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Discriminator(nn.Module):\n    def __init__(self, hid_channels, in_channels=3):\n        super().__init__()\n        self.block = nn.Sequential(\n            Downsampling(in_channels, hid_channels, norm=False), # 64x128x128\n            Downsampling(hid_channels, hid_channels*2), # 128x64x64\n            Downsampling(hid_channels*2, hid_channels*4), # 256x32x32\n            Downsampling(hid_channels*4, hid_channels*8, stride=1), # 512x31x31\n            nn.Conv2d(hid_channels*8, 1, kernel_size=4, padding=1), # 1x30x30\n        )\n        \n    def forward(self, x):\n        return self.block(x)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class ImageBuffer(object):\n    def __init__(self, buffer_size):\n        self.buffer_size = buffer_size\n        if self.buffer_size > 0:\n            # the current capacity of the buffer\n            self.curr_cap = 0\n            # initialize buffer as empty list\n            self.buffer = []\n    \n    def __call__(self, imgs):\n        # the buffer is not used\n        if self.buffer_size == 0:\n            return imgs\n        \n        return_imgs = []\n        for img in imgs:\n            img = img.unsqueeze(dim=0)\n            \n            # fill buffer to maximum capacity\n            if self.curr_cap < self.buffer_size:\n                self.curr_cap += 1\n                self.buffer.append(img)\n                return_imgs.append(img)\n            else:\n                p = np.random.uniform(low=0., high=1.)\n                \n                # swap images between input and buffer with probability 0.5\n                if p > 0.5:\n                    idx = np.random.randint(low=0, high=self.buffer_size)\n                    tmp = self.buffer[idx].clone()\n                    self.buffer[idx] = img\n                    return_imgs.append(tmp)\n                else:\n                    return_imgs.append(img)\n        return torch.cat(return_imgs, dim=0)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"MODEL_CONFIG = {\n    # the type of generator, and the number of residual blocks if ResNet generator is used\n    \"gen_name\": \"unet\", # types: 'unet', 'resnet'\n    \"num_resblocks\": 6,\n    # the number of filters in the first layer for the generators and discriminators\n    \"hid_channels\": 64,\n    # using DeepSpeed's FusedAdam (currently GPU only) is slightly faster\n    \"optimizer\": ds.ops.adam.FusedAdam if torch.cuda.is_available() else torch.optim.Adam,\n    # the learning rate and beta parameters for the Adam optimizer\n    \"lr\": 2e-4,\n    \"betas\": (0.5, 0.999),\n    # the weights used in the identity loss and cycle loss\n    \"lambda_idt\": 0.5,\n    \"lambda_cycle\": (10, 10), # (MPM direction, PMP direction)\n    # the size of the buffer that stores previously generated images\n    \"buffer_size\": 100,\n    # the number of epochs for training\n    \"num_epochs\": 18 if not DEBUG else 2,\n    # the number of epochs before starting the learning rate decay\n    \"decay_epochs\": 18 if not DEBUG else 1,\n}","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CycleGAN(L.LightningModule):\n    def __init__(\n        self,\n        gen_name,\n        num_resblocks,\n        hid_channels,\n        optimizer,\n        lr,\n        betas,\n        lambda_idt,\n        lambda_cycle,\n        buffer_size,\n        num_epochs,\n        decay_epochs,\n    ):\n        super().__init__()\n        self.save_hyperparameters(ignore=[\"optimizer\"])\n        self.optimizer = optimizer\n        self.automatic_optimization = False\n        \n        # define generators and discriminators\n        self.gen_PM = get_gen(gen_name, hid_channels, num_resblocks)\n        self.gen_MP = get_gen(gen_name, hid_channels, num_resblocks)\n        self.disc_M = Discriminator(hid_channels)\n        self.disc_P = Discriminator(hid_channels)\n        \n        # initialize buffers to store fake images\n        self.buffer_fake_M = ImageBuffer(buffer_size)\n        self.buffer_fake_P = ImageBuffer(buffer_size)\n        \n    def forward(self, img):\n        return self.gen_PM(img)   \n            \n    def init_weights(self):\n        def init_fn(m):\n            if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d, nn.InstanceNorm2d)):\n                nn.init.normal_(m.weight, 0.0, 0.02)\n                if m.bias is not None:\n                    nn.init.constant_(m.bias, 0.0)\n        \n        for net in [self.gen_PM, self.gen_MP, self.disc_M, self.disc_P]:\n            net.apply(init_fn)\n        \n    def setup(self, stage):\n        if stage == \"fit\":\n            self.init_weights()\n            print(\"Model initialized.\")\n            \n    def get_lr_scheduler(self, optimizer):\n        def lr_lambda(epoch):\n            len_decay_phase = self.hparams.num_epochs - self.hparams.decay_epochs + 1.0\n            curr_decay_step = max(0, epoch - self.hparams.decay_epochs + 1.0)\n            val = 1.0 - curr_decay_step / len_decay_phase\n            return max(0.0, val)\n        \n        return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lr_lambda)\n    \n    def configure_optimizers(self):\n        opt_config = {\n            \"lr\": self.hparams.lr,\n            \"betas\": self.hparams.betas,\n        }\n        opt_gen = self.optimizer(\n            list(self.gen_PM.parameters()) + list(self.gen_MP.parameters()),\n            **opt_config,\n        )\n        opt_disc = self.optimizer(\n            list(self.disc_M.parameters()) + list(self.disc_P.parameters()),\n            **opt_config,\n        )\n        optimizers = [opt_gen, opt_disc]\n        schedulers = [self.get_lr_scheduler(opt) for opt in optimizers]\n        return optimizers, schedulers\n        \n    def adv_criterion(self, y_hat, y):\n        return F.mse_loss(y_hat, y)\n    \n    def recon_criterion(self, y_hat, y):\n        return F.l1_loss(y_hat, y)\n    \n    def get_adv_loss(self, fake, disc):\n        fake_hat = disc(fake)\n        real_labels = torch.ones_like(fake_hat)\n        adv_loss = self.adv_criterion(fake_hat, real_labels)\n        return adv_loss\n    \n    def get_idt_loss(self, real, idt, lambda_cycle):\n        idt_loss = self.recon_criterion(idt, real)\n        return self.hparams.lambda_idt * lambda_cycle * idt_loss\n    \n    def get_cycle_loss(self, real, recon, lambda_cycle):\n        cycle_loss = self.recon_criterion(recon, real)\n        return lambda_cycle * cycle_loss\n    \n    def get_gen_loss(self):\n        # calculate adversarial loss\n        adv_loss_PM = self.get_adv_loss(self.fake_M, self.disc_M)\n        adv_loss_MP = self.get_adv_loss(self.fake_P, self.disc_P)\n        total_adv_loss = adv_loss_PM + adv_loss_MP\n        \n        # calculate identity loss\n        lambda_cycle = self.hparams.lambda_cycle\n        idt_loss_MM = self.get_idt_loss(self.real_M, self.idt_M, lambda_cycle[0])\n        idt_loss_PP = self.get_idt_loss(self.real_P, self.idt_P, lambda_cycle[1])\n        total_idt_loss = idt_loss_MM + idt_loss_PP\n        \n        # calculate cycle loss\n        cycle_loss_MPM = self.get_cycle_loss(self.real_M, self.recon_M, lambda_cycle[0])\n        cycle_loss_PMP = self.get_cycle_loss(self.real_P, self.recon_P, lambda_cycle[1])\n        total_cycle_loss = cycle_loss_MPM + cycle_loss_PMP\n        \n        # combine losses\n        gen_loss = total_adv_loss + total_idt_loss + total_cycle_loss\n        return gen_loss\n    \n    def get_disc_loss(self, real, fake, disc):\n        # calculate loss on real images\n        real_hat = disc(real)\n        real_labels = torch.ones_like(real_hat)\n        real_loss = self.adv_criterion(real_hat, real_labels)\n        \n        # calculate loss on fake images\n        fake_hat = disc(fake.detach())\n        fake_labels = torch.zeros_like(fake_hat)\n        fake_loss = self.adv_criterion(fake_hat, fake_labels)\n        \n        # combine losses\n        disc_loss = (fake_loss + real_loss) * 0.5\n        return disc_loss\n    \n    def get_disc_loss_M(self):\n        fake_M = self.buffer_fake_M(self.fake_M)\n        return self.get_disc_loss(self.real_M, fake_M, self.disc_M)\n    \n    def get_disc_loss_P(self):\n        fake_P = self.buffer_fake_P(self.fake_P)\n        return self.get_disc_loss(self.real_P, fake_P, self.disc_P)\n    \n    def training_step(self, batch, batch_idx):\n        self.real_M = batch[\"monet\"]\n        self.real_P = batch[\"photo\"]\n        opt_gen, opt_disc = self.optimizers()\n\n        # generate fake images\n        self.fake_M = self.gen_PM(self.real_P)\n        self.fake_P = self.gen_MP(self.real_M)\n        \n        # generate identity images\n        self.idt_M = self.gen_PM(self.real_M)\n        self.idt_P = self.gen_MP(self.real_P)\n        \n        # reconstruct images\n        self.recon_M = self.gen_PM(self.fake_P)\n        self.recon_P = self.gen_MP(self.fake_M)\n    \n        # train generators\n        self.toggle_optimizer(opt_gen)\n        gen_loss = self.get_gen_loss()        \n        opt_gen.zero_grad()\n        self.manual_backward(gen_loss)\n        opt_gen.step()\n        self.untoggle_optimizer(opt_gen)\n        \n        # train discriminators\n        self.toggle_optimizer(opt_disc)\n        disc_loss_M = self.get_disc_loss_M()\n        disc_loss_P = self.get_disc_loss_P()\n        opt_disc.zero_grad()\n        self.manual_backward(disc_loss_M)\n        self.manual_backward(disc_loss_P)\n        opt_disc.step()\n        self.untoggle_optimizer(opt_disc)\n        \n        # record training losses\n        metrics = {\n            \"gen_loss\": gen_loss,\n            \"disc_loss_M\": disc_loss_M,\n            \"disc_loss_P\": disc_loss_P,\n        }\n        self.log_dict(metrics, on_step=False, on_epoch=True, prog_bar=True)\n        \n    def validation_step(self, batch, batch_idx):\n        self.display_results(batch, batch_idx, \"validate\")\n    \n    def test_step(self, batch, batch_idx):\n        self.display_results(batch, batch_idx, \"test\")\n        \n    def predict_step(self, batch, batch_idx):\n        return self(batch)\n    \n    def display_results(self, batch, batch_idx, stage):\n        real_P = batch\n        fake_M = self(real_P)\n        \n        if stage == \"validate\":\n            title = f\"Epoch {self.current_epoch+1}: Photo-to-Monet Translation\"\n        else:\n            title = f\"Sample {batch_idx+1}: Photo-to-Monet Translation\"\n\n        show_img(\n            torch.cat([real_P, fake_M], dim=0),\n            nrow=len(real_P),\n            title=title,\n        )\n    \n    def on_train_epoch_start(self):\n        # record learning rates\n        curr_lr = self.lr_schedulers()[0].get_last_lr()[0]\n        self.log(\"lr\", curr_lr, on_step=False, on_epoch=True, prog_bar=True)\n        \n    def on_train_epoch_end(self):\n        # update learning rates\n        for sch in self.lr_schedulers():\n            sch.step()\n        \n        # print current state of epoch\n        logged_values = self.trainer.progress_bar_metrics\n        print(\n            f\"Epoch {self.current_epoch+1}\",\n            *[f\"{k}: {v:.5f}\" for k, v in logged_values.items()],\n            sep=\" - \",\n        )\n        \n    def on_train_end(self):\n        print(\"Training ended.\")\n        \n    def on_predict_epoch_end(self):\n        predictions = self.trainer.predict_loop.predictions\n        num_batches = len(predictions)\n        batch_size = predictions[0].shape[0]\n        last_batch_diff = batch_size - predictions[-1].shape[0]\n        print(f\"Number of images generated: {num_batches*batch_size-last_batch_diff}.\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"TRAIN_CONFIG = {\n    \"accelerator\": \"gpu\" if torch.cuda.is_available() else \"cpu\",\n    \n    # train on 16-bit precision\n    \"precision\": \"16-mixed\" if torch.cuda.is_available() else 32,\n    \n    # train on single GPU\n    \"devices\": 1,\n    \n    # save checkpoint only for last epoch by default\n    \"enable_checkpointing\": True,\n    \n    # disable logging for simplicity\n    \"logger\": False,\n    \n    # the number of epochs for training (we limit the number of train/predict batches during debugging)\n    \"max_epochs\": MODEL_CONFIG[\"num_epochs\"],\n    \"limit_train_batches\": 1.0 if not DEBUG else 2,\n    \"limit_predict_batches\": 1.0 if not DEBUG else 5,\n    \n    # the maximum amount of time for training, in case we exceed run-time of 5 hours\n    \"max_time\": {\"hours\": 4, \"minutes\": 55},\n    \n    # use a small subset of photos for validation/testing (we limit here for flexibility)\n    \"limit_val_batches\": 1,\n    \"limit_test_batches\": 5,\n    \n    # disable sanity check before starting the training routine\n    \"num_sanity_val_steps\": 0,\n    \n    # the frequency to visualize the progress of adding Monet style\n    \"check_val_every_n_epoch\": 6 if not DEBUG else 1,\n}","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"dm = CustomDataModule(**DM_CONFIG)\nmodel = CycleGAN(**MODEL_CONFIG)\ntrainer = L.Trainer(**TRAIN_CONFIG)\ntrainer.fit(model, datamodule=dm)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"_ = trainer.test(model, datamodule=dm)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"predictions = trainer.predict(model, datamodule=dm)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"os.makedirs(\"../images\", exist_ok=True)\nidx = 0\nfor tensor in predictions:\n    for monet in tensor:\n        save_image(\n            monet.float().squeeze() * 0.5 + 0.5, \n            fp=f\"../images/{idx}.jpg\",\n        )\n        idx += 1\n\nshutil.make_archive(\"/kaggle/working/images\", \"zip\", \"/kaggle/images\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}