{"metadata":{"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":21755,"databundleVersionId":1475600,"sourceType":"competition"}],"dockerImageVersionId":30747,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true},"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.10.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nimport itertools\n\nimport numpy as np\n\nimport torch \nimport torch.nn as nn\nimport torch.nn.functional as F\n\nfrom torch.optim import Adam\nfrom torch.optim.lr_scheduler import StepLR\nfrom torch.utils.data import Dataset, DataLoader\n\nimport torchvision.transforms.v2 as T\nimport torchvision.transforms.v2.functional as TF\nfrom torchvision.utils import make_grid\n\nimport math\n\nimport random\n\nimport functools\n\nfrom tqdm import tqdm\n\nimport warnings\nwarnings.filterwarnings('ignore')","metadata":{"execution":{"iopub.status.busy":"2024-09-04T07:59:23.208307Z","iopub.execute_input":"2024-09-04T07:59:23.208577Z","iopub.status.idle":"2024-09-04T07:59:28.643390Z","shell.execute_reply.started":"2024-09-04T07:59:23.208553Z","shell.execute_reply":"2024-09-04T07:59:28.642226Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Creating datasets and dataloaders","metadata":{}},{"cell_type":"code","source":"torch.cuda.empty_cache()","metadata":{"execution":{"iopub.status.busy":"2024-09-04T07:59:28.644929Z","iopub.execute_input":"2024-09-04T07:59:28.645330Z","iopub.status.idle":"2024-09-04T07:59:28.649550Z","shell.execute_reply.started":"2024-09-04T07:59:28.645304Z","shell.execute_reply":"2024-09-04T07:59:28.648411Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from glob import glob\nfrom PIL import Image","metadata":{"execution":{"iopub.status.busy":"2024-09-04T07:59:28.650827Z","iopub.execute_input":"2024-09-04T07:59:28.651099Z","iopub.status.idle":"2024-09-04T07:59:28.659439Z","shell.execute_reply.started":"2024-09-04T07:59:28.651077Z","shell.execute_reply":"2024-09-04T07:59:28.658586Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CustomDataset(Dataset):\n    def __init__(self, path, transform=None):\n        super().__init__()\n        \n        self.path = path\n        self.files = glob(f'{path}/*.jpg')\n        self.transform = transform\n        \n    def __len__(self):\n        return len(self.files) \n    \n    def __getitem__(self, idx):\n        if torch.is_tensor(idx):\n            idx = idx.tolist()\n        \n        img = Image.open(self.files[idx])\n        \n        if self.transform:\n            img = self.transform(img)\n            \n        return img","metadata":{"execution":{"iopub.status.busy":"2024-09-04T07:59:28.661746Z","iopub.execute_input":"2024-09-04T07:59:28.662046Z","iopub.status.idle":"2024-09-04T07:59:28.670343Z","shell.execute_reply.started":"2024-09-04T07:59:28.662024Z","shell.execute_reply":"2024-09-04T07:59:28.669432Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"transform = T.Compose(\n    [\n        T.Resize((268,268)),\n        T.RandomCrop((256,256)),\n        T.RandomHorizontalFlip(),\n        T.ColorJitter(contrast=0.2, saturation=0.2, hue=0.1),\n        T.ToTensor(),\n        T.Normalize([0.5,0.5,0.5],[0.5,0.5,0.5])\n    ]\n)","metadata":{"execution":{"iopub.status.busy":"2024-09-04T07:59:28.671523Z","iopub.execute_input":"2024-09-04T07:59:28.671776Z","iopub.status.idle":"2024-09-04T07:59:28.680388Z","shell.execute_reply.started":"2024-09-04T07:59:28.671755Z","shell.execute_reply":"2024-09-04T07:59:28.679620Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def inverse_transform(img):\n    return (img * 0.5) + 0.5","metadata":{"execution":{"iopub.status.busy":"2024-09-04T07:59:28.681623Z","iopub.execute_input":"2024-09-04T07:59:28.682035Z","iopub.status.idle":"2024-09-04T07:59:28.689810Z","shell.execute_reply.started":"2024-09-04T07:59:28.682001Z","shell.execute_reply":"2024-09-04T07:59:28.688951Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"monet_ds = CustomDataset('/kaggle/input/gan-getting-started/monet_jpg', transform=transform)\nphoto_ds = CustomDataset('/kaggle/input/gan-getting-started/photo_jpg', transform=transform)","metadata":{"execution":{"iopub.status.busy":"2024-09-04T07:59:28.691211Z","iopub.execute_input":"2024-09-04T07:59:28.691519Z","iopub.status.idle":"2024-09-04T07:59:29.173344Z","shell.execute_reply.started":"2024-09-04T07:59:28.691491Z","shell.execute_reply":"2024-09-04T07:59:29.172371Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CombinedLoader:\n    def __init__(self, photo_ds, monet_ds, batch_size, num_workers):\n        self.photo_ds = photo_ds\n        self.monet_ds = monet_ds\n        self.batch_size = batch_size\n        self.num_workers = num_workers\n        self.len = max(len(photo_ds), len(monet_ds))//batch_size\n\n    def __iter__(self):\n        self.photo_loader = DataLoader(self.photo_ds, batch_size=self.batch_size, drop_last=True, num_workers=self.num_workers, shuffle=True)\n        self.monet_loader = DataLoader(self.monet_ds, batch_size=self.batch_size, drop_last=True, num_workers=self.num_workers, shuffle=True)\n        self.photo_iter = itertools.cycle(self.photo_loader)\n        self.monet_iter = itertools.cycle(self.monet_loader)\n        self.counter = 0\n        return self\n    \n    def __len__(self):\n        return self.len\n\n    def __next__(self):\n        if self.counter > self.len:\n            raise StopIteration\n        \n        self.counter += 1\n        return next(self.photo_iter), next(self.monet_iter)     ","metadata":{"execution":{"iopub.status.busy":"2024-09-04T07:59:29.174577Z","iopub.execute_input":"2024-09-04T07:59:29.175253Z","iopub.status.idle":"2024-09-04T07:59:29.183844Z","shell.execute_reply.started":"2024-09-04T07:59:29.175219Z","shell.execute_reply":"2024-09-04T07:59:29.182870Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Defining the Model\noriginal paper - [Unpaired Image-to-Image Translation using Cycle-Consistent Adversarial Networks](https://arxiv.org/abs/1703.10593)\n\n\n### The Generator\nThe original paper proposes a generator architecture made up of multiple residual blocks, however that fails to converge easily.\nThe U-Net architecture that is used by pix2pix converges much faster\n\nReflection padding is used to reduce artifacts\n\nIn the upsample blocks, we don't use Conv2dTranspose. Instead, we use a NearestNeighbour Upsample block followed by a Conv2d layer.\nThis is done to further reduce artifacts and checkboard patterns in the generated images. Source: https://distill.pub/2016/deconv-checkerboard/\n\n### The Discriminator\nWe use a PatchGAN discriminator as mentioned in the original paper","metadata":{}},{"cell_type":"code","source":"class G_down_block(nn.Module):\n    def __init__(self, in_channels, out_channels, norm=True):\n        super().__init__()\n        self.layers = nn.Sequential(\n            nn.Conv2d(\n                in_channels, \n                out_channels, \n                kernel_size=4, \n                stride=2, \n                padding=1, \n                padding_mode=\"reflect\"\n            )\n        )\n        \n        if norm:\n            self.layers.append(nn.InstanceNorm2d(out_channels))\n        \n        self.layers.append(nn.LeakyReLU(0.2))\n        \n    def forward(self, x):\n        return self.layers(x) \n\nclass G_up_block(nn.Module):\n    def __init__(self, in_channels, out_channels, dropout=False, act=\"leaky-relu\"):\n        super().__init__()\n        self.layers = nn.Sequential(\n            nn.UpsamplingNearest2d(scale_factor=2),\n            nn.Conv2d(\n                in_channels, \n                out_channels, \n                kernel_size=3, \n                padding=1,\n                padding_mode=\"reflect\"\n            )\n        )\n        \n        self.layers.append(nn.InstanceNorm2d(out_channels))\n        \n        if act == \"tanh\":\n            self.layers.append(nn.Tanh())\n        else:\n            self.layers.append(nn.LeakyReLU(0.2))\n        \n        if dropout:\n            self.layers.append(nn.Dropout(0.5))\n        \n    def forward(self, x):\n        return self.layers(x) \n\nclass Generator(nn.Module):\n    def __init__(self):\n        super().__init__()\n        self.down = nn.Sequential(\n            G_down_block(3,64,norm=False),\n            G_down_block(64,128),\n            G_down_block(128,256),\n            G_down_block(256,512),\n            G_down_block(512,512),\n            G_down_block(512,512),\n            G_down_block(512,512),\n            G_down_block(512,512,norm=False)\n        )\n        self.up = nn.Sequential(\n            G_up_block(2*512,512,dropout=True),\n            G_up_block(2*512,512,dropout=True),\n            G_up_block(2*512,512,dropout=True),\n            G_up_block(2*512,512),\n            G_up_block(2*512,256),\n            G_up_block(2*256,128),\n            G_up_block(2*128,64),\n            G_up_block(2*64,3,act=\"tanh\"),\n        )\n    def forward(self, x):\n        buffer = []\n        for net in self.down:\n            x = net(x)\n            buffer.append(x)\n        \n        for res, net in zip(buffer[::-1], self.up):\n            x = net(torch.cat([x,res], dim=1))\n            \n        return x\n\nclass Discriminator(nn.Module):\n    def __init__(self, in_channels=3):\n        super().__init__()\n        self.layers = nn.Sequential(\n            nn.Conv2d(in_channels,64,4,2,1),\n            nn.LeakyReLU(0.2,True),\n            nn.Conv2d(64,128,4,2,1,bias=True),\n            nn.InstanceNorm2d(128),\n            nn.LeakyReLU(0.2,True),\n            nn.Conv2d(128,256,4,2,1,bias=True),\n            nn.InstanceNorm2d(256),\n            nn.LeakyReLU(0.2,True),\n            nn.Conv2d(256,512,4,1,1,bias=True),\n            nn.InstanceNorm2d(512),\n            nn.LeakyReLU(0.2,True),\n            nn.Conv2d(512,1,4,1,1),\n        )\n    \n    def forward(self, x):\n        return self.layers(x)","metadata":{"execution":{"iopub.status.busy":"2024-09-04T07:59:29.185230Z","iopub.execute_input":"2024-09-04T07:59:29.185491Z","iopub.status.idle":"2024-09-04T07:59:29.206882Z","shell.execute_reply.started":"2024-09-04T07:59:29.185469Z","shell.execute_reply":"2024-09-04T07:59:29.206054Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CycleGAN(nn.Module):\n    def __init__(self):\n        super().__init__()\n        self.photo_G = Generator()\n        self.photo_D = Discriminator()\n        \n        self.monet_G = Generator()\n        self.monet_D = Discriminator()\n\n        self.apply(self._init_weights)\n\n        print(f'Cycle GAN initialized with {self._get_num_params()} parameters')\n\n    def _init_weights(self, module):\n        if isinstance(module, (nn.Conv2d, nn.ConvTranspose2d)):\n            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)\n            if module.bias is not None:\n                torch.nn.init.zeros_(module.bias)\n        \n    def _get_num_params(self):\n        p_G = sum([p.numel() for p in self.photo_G.parameters()])\n        p_D = sum([p.numel() for p in self.photo_D.parameters()])\n        m_G = sum([p.numel() for p in self.monet_G.parameters()])\n        m_D = sum([p.numel() for p in self.monet_D.parameters()])\n        \n        return p_G + p_D + m_G + m_D\n    \n    def photo_to_monet(self, photo):\n        self.monet_G.eval()\n        with torch.no_grad():\n            out = self.monet_G(photo)\n        return out\n    \n    def monet_to_photo(self, monet):\n        self.photo_G.eval()\n        with torch.no_grad():\n            out = self.photo_G(monet)\n        return out","metadata":{"execution":{"iopub.status.busy":"2024-09-04T07:59:29.210348Z","iopub.execute_input":"2024-09-04T07:59:29.210631Z","iopub.status.idle":"2024-09-04T07:59:29.222412Z","shell.execute_reply.started":"2024-09-04T07:59:29.210602Z","shell.execute_reply":"2024-09-04T07:59:29.221334Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Image Buffer\nInspired by the original paper.\nThis is used in training to randomly fetch images for training the discriminator ","metadata":{}},{"cell_type":"code","source":"class ImageBuffer:\n    def __init__(self, size=70):\n        self.size = size\n        self.buffer = []\n    def pass_images(self, imgs):\n        out_buffer = []\n        for img in imgs:\n            if len(self.buffer) < self.size:\n                self.buffer.append(img)\n                out_buffer.append(img)\n            else:\n                #faster than random.choice https://stackoverflow.com/questions/6824681/get-a-random-boolean-in-python\n                if bool(random.getrandbits(1)): \n                    fetch_img_idx = random.choice(range(self.size))\n                    out_buffer.append(self.buffer[fetch_img_idx])\n                    self.buffer[fetch_img_idx] = img\n                else:\n                    out_buffer.append(img)\n        return torch.stack(out_buffer, dim=0)","metadata":{"execution":{"iopub.status.busy":"2024-09-04T07:59:29.223441Z","iopub.execute_input":"2024-09-04T07:59:29.223695Z","iopub.status.idle":"2024-09-04T07:59:29.235635Z","shell.execute_reply.started":"2024-09-04T07:59:29.223674Z","shell.execute_reply":"2024-09-04T07:59:29.234705Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# The Trainer\n\n### Logging\nWe use tensorboard for logging (disabled for submission notebooks)\nScalar metrics are logged for every step and Images are logged at the end of every epoch \n\n### Learning Rate Scheduler\nA custom lr scheduler is used. (for more control and also pytorch lr schedulers confuse me)","metadata":{}},{"cell_type":"code","source":"from dataclasses import dataclass\nfrom typing import Tuple\n\nfrom torch.utils.tensorboard import SummaryWriter\n\nimport itertools","metadata":{"execution":{"iopub.status.busy":"2024-09-04T07:59:29.236740Z","iopub.execute_input":"2024-09-04T07:59:29.237329Z","iopub.status.idle":"2024-09-04T07:59:41.181953Z","shell.execute_reply.started":"2024-09-04T07:59:29.237298Z","shell.execute_reply":"2024-09-04T07:59:41.181135Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"@dataclass\nclass TrainerConfig:\n    log_dir: str\n    batch_size: int = 16\n    num_workers: int = 2\n    device: str|None = None\n\n    h_lambda: float = 10\n    lr: float = 3e-4\n    betas: Tuple[float, float] = (0.9, 0.999)\n    eps: float = 1e-8\n        \n    grad_norm: float|None = None\n        \n    weight_decay: float = 0\n    warmup_iters: int = 500\n    min_lr: float = 1e-5\n    lr_decay_iters: int = 4000\n\n    def __post_init__(self):\n        if self.device is None:\n            self.device = \"cuda\" if torch.cuda.is_available() else \"cpu\"","metadata":{"execution":{"iopub.status.busy":"2024-09-04T07:59:41.183022Z","iopub.execute_input":"2024-09-04T07:59:41.183533Z","iopub.status.idle":"2024-09-04T07:59:41.192116Z","shell.execute_reply.started":"2024-09-04T07:59:41.183501Z","shell.execute_reply":"2024-09-04T07:59:41.191224Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Trainer:\n    def __init__(self, config, model):\n        self.model = model\n        self.device = config.device\n        \n        self.lr = config.lr\n        \n        self.grad_norm = config.grad_norm\n        \n        self.warmup_iters = config.warmup_iters\n        self.min_lr = config.min_lr\n        self.lr_decay_iters = config.lr_decay_iters\n\n        self.h_lambda = config.h_lambda\n        \n        self.g_optim, self.d_optim = self._get_optimizers(config)\n        \n        #self.log_writer = SummaryWriter(config.log_dir)\n        \n        self.loader = self._get_dataloaders(config)\n        \n        self.fake_monet_buffer = ImageBuffer()\n        self.fake_photo_buffer = ImageBuffer()\n        \n        self.metrics_buffer = []\n\n    #custom lr scheduler inspired by https://github.com/karpathy/nanoGPT\n    def get_lr(self, it):\n        # 1) linear warmup for warmup_iters steps\n        if it < self.warmup_iters:\n            return self.lr * it / self.warmup_iters\n        # 2) if it > lr_decay_iters, return min learning rate\n        if it > self.lr_decay_iters:\n            return self.min_lr\n        # 3) in between, use cosine decay down to min learning rate\n        decay_ratio = (it - self.warmup_iters) / (self.lr_decay_iters - self.warmup_iters)\n        assert 0 <= decay_ratio <= 1\n        coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio)) # coeff ranges 0..1\n        return self.min_lr + coeff * (self.lr - self.min_lr)\n        \n    def _get_optimizers(self, config):\n        g_optim = Adam(\n            itertools.chain(\n                self.model.photo_G.parameters(), \n                self.model.monet_G.parameters()\n            ),\n            lr=config.lr,\n            betas=config.betas,\n            eps=config.eps,\n            weight_decay=config.weight_decay\n        )\n        d_optim = Adam(\n            itertools.chain(\n                self.model.photo_D.parameters(), \n                self.model.monet_D.parameters()\n            ),\n            lr=config.lr,\n            betas=config.betas,\n            eps=config.eps,\n            weight_decay=config.weight_decay\n        )\n        return g_optim, d_optim\n    \n    def _get_dataloaders(self, config):\n        return CombinedLoader(photo_ds, monet_ds, batch_size=config.batch_size, num_workers=config.num_workers)\n       \n    def _set_discriminator_requires_grad(self, requires_grad=False):\n        for net in [self.model.photo_D, self.model.monet_D]:\n            if net is not None:\n                for param in net.parameters():\n                    param.requires_grad = requires_grad\n        \n    def _get_gan_loss(self, x, is_real):\n        target = torch.full(x.shape, float(is_real), device=self.device)\n        return F.mse_loss(x, target)\n    \n    def _get_discriminator_loss(self, net, real, fake):\n        loss_real = self._get_gan_loss(net(real), is_real=True)\n        loss_fake = self._get_gan_loss(net(fake), is_real=False)\n        loss = (loss_real + loss_fake) * 0.5\n        loss.backward()\n        return loss\n    \n    def train_step(self, real_photo, real_monet):\n        self.model.train()\n        \n        fake_photo = self.model.photo_G(real_monet)\n        fake_monet = self.model.monet_G(real_photo)\n        cycle_photo = self.model.photo_G(fake_monet)\n        cycle_monet = self.model.monet_G(fake_photo)\n        \n        #lock discriminators\n        self._set_discriminator_requires_grad(False)\n        \n        #generator optimization\n        self.g_optim.zero_grad()\n        \n        #identity loss\n        id_loss_1 = F.l1_loss(self.model.photo_G(real_photo), real_photo) * 0.5 * self.h_lambda\n        id_loss_2 = F.l1_loss(self.model.monet_G(real_monet), real_monet) * 0.5 * self.h_lambda\n        id_loss = id_loss_1 + id_loss_2\n        \n        #generator loss\n        photo_G_loss = self._get_gan_loss(\n            self.model.photo_D(fake_photo),\n            is_real=True\n        )\n        monet_G_loss = self._get_gan_loss(\n            self.model.monet_D(fake_monet),\n            is_real=True\n        )\n        gen_loss = photo_G_loss + monet_G_loss\n        \n        #cycle loss\n        loss_photo_cycle = F.l1_loss(cycle_photo, real_photo) * self.h_lambda\n        loss_monet_cycle = F.l1_loss(cycle_monet, real_monet) * self.h_lambda\n        cycle_loss = loss_photo_cycle + loss_monet_cycle\n        \n        #backprop for generators\n        total_G_loss =  id_loss + gen_loss + cycle_loss\n        total_G_loss.backward()\n        \n        nn.utils.clip_grad_norm_(\n            itertools.chain(\n                self.model.photo_G.parameters(), \n                self.model.monet_G.parameters()\n            ), \n            1.0\n        )\n        \n        self.g_optim.step()\n        \n        #unfreeze discriminators\n        self._set_discriminator_requires_grad(True)\n        \n        #discriminator optimization\n        self.d_optim.zero_grad()\n        \n        photo_D_loss = self._get_discriminator_loss(\n            self.model.photo_D,\n            real_photo,\n            self.fake_photo_buffer.pass_images(fake_photo.detach())\n        )\n        monet_D_loss = self._get_discriminator_loss(\n            self.model.monet_D,\n            real_monet,\n            self.fake_monet_buffer.pass_images(fake_monet.detach())\n        )\n        \n        if self.grad_norm:\n            nn.utils.clip_grad_norm_(\n                itertools.chain(\n                    self.model.photo_D.parameters(), \n                    self.model.monet_D.parameters()\n                ), \n                self.grad_norm\n            )\n\n        self.d_optim.step()\n\n        return {\n            \"monet_g_loss\": monet_G_loss.item(),\n            \"photo_g_loss\": photo_G_loss.item(),\n            \"monet_d_loss\": monet_D_loss.item(),\n            \"photo_d_loss\": photo_D_loss.item(),\n            \"identity_loss\": id_loss.item(),\n            \"cyclic_loss\": cycle_loss.item(),\n            \"total_g_loss\": total_G_loss.item(),\n            \"total_d_loss\": (photo_D_loss + monet_D_loss).item()\n        }\n    \n    def eval_step(self, step):\n        eval_photo = random.choice(photo_ds).to(self.device)\n        gen_monet = self.model.photo_to_monet(eval_photo.unsqueeze(0)).squeeze(0)\n        \n        photo2monet_grid = inverse_transform(make_grid([eval_photo, gen_monet]))\n\n        #self.log_writer.add_image('photo_to_monet', photo2monet_grid, step)\n        display(TF.to_pil_image(photo2monet_grid))\n        \n        \n    \n    def train(self, epochs):\n        step = 0\n        tqdm_g_loss = 0.0  \n        tqdm_d_loss = 0.0\n        for epoch in range(epochs):\n            running_total_g_loss = 0.0\n            running_total_d_loss = 0.0\n            for photo, monet in (pbar := tqdm(self.loader)):\n                pbar.set_description(f\"Epoch{epoch:2d} G_loss:{tqdm_g_loss:5.3f} D_loss:{tqdm_d_loss:5.3f}\")\n                \n                photo, monet = photo.to(self.device), monet.to(self.device)\n                \n                new_lr = self.get_lr(step)\n                for param_group in itertools.chain(\n                    self.d_optim.param_groups, \n                    self.g_optim.param_groups\n                ):\n                    param_group['lr'] = new_lr\n\n                metrics_dict = self.train_step(photo, monet)\n                self.metrics_buffer.append(metrics_dict)\n                \n                tqdm_g_loss = metrics_dict[\"total_g_loss\"]\n                tqdm_d_loss = metrics_dict[\"total_d_loss\"]\n\n                running_total_g_loss += tqdm_g_loss\n                running_total_d_loss += tqdm_d_loss\n\n                # for metric in metrics_dict.keys():\n                #     self.log_writer.add_scalar(str(metric), metrics_dict[metric], step)\n                \n                step += 1\n\n            # cli log \n            print(f\"g_loss: {running_total_g_loss/len(self.loader)} d_loss: {running_total_d_loss/len(self.loader)}\")\n\n            # image log\n            self.eval_step(step)\n","metadata":{"execution":{"iopub.status.busy":"2024-09-04T07:59:41.193403Z","iopub.execute_input":"2024-09-04T07:59:41.193731Z","iopub.status.idle":"2024-09-04T07:59:41.239991Z","shell.execute_reply.started":"2024-09-04T07:59:41.193701Z","shell.execute_reply":"2024-09-04T07:59:41.239156Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Training","metadata":{}},{"cell_type":"code","source":"config = TrainerConfig(\n    log_dir=\"./logs\",\n    batch_size=8,\n    lr=2e-4, #lr used in the paper\n    grad_norm=1.0,\n    weight_decay=0.01\n)\n\nmodel = CycleGAN().to(config.device)\n\ntrainer = Trainer(\n    config,\n    model\n)","metadata":{"execution":{"iopub.status.busy":"2024-09-04T07:59:41.241112Z","iopub.execute_input":"2024-09-04T07:59:41.241465Z","iopub.status.idle":"2024-09-04T07:59:43.148942Z","shell.execute_reply.started":"2024-09-04T07:59:41.241435Z","shell.execute_reply":"2024-09-04T07:59:43.148190Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"trainer.train(3)","metadata":{"execution":{"iopub.status.busy":"2024-09-04T07:59:43.149979Z","iopub.execute_input":"2024-09-04T07:59:43.150261Z","iopub.status.idle":"2024-09-04T08:34:06.364125Z","shell.execute_reply.started":"2024-09-04T07:59:43.150237Z","shell.execute_reply":"2024-09-04T08:34:06.363086Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"photo_test_ds = CustomDataset('/kaggle/input/gan-getting-started/photo_jpg', transform=T.Compose([T.ToTensor(), T.Normalize([0.5,0.5,0.5], [0.5,0.5,0.5])]))\ntest_loader = DataLoader(photo_test_ds, num_workers=config.num_workers, drop_last=False, batch_size=config.batch_size)","metadata":{"execution":{"iopub.status.busy":"2024-09-04T08:34:06.365667Z","iopub.execute_input":"2024-09-04T08:34:06.365978Z","iopub.status.idle":"2024-09-04T08:34:06.396536Z","shell.execute_reply.started":"2024-09-04T08:34:06.365948Z","shell.execute_reply":"2024-09-04T08:34:06.395696Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Save model weights (optional)","metadata":{}},{"cell_type":"code","source":"#torch.save(trainer.model.state_dict(), \"/tmp/7_epochs_UNET_modified.pt\")","metadata":{"execution":{"iopub.status.busy":"2024-09-04T08:34:06.397679Z","iopub.execute_input":"2024-09-04T08:34:06.399181Z","iopub.status.idle":"2024-09-04T08:34:06.403223Z","shell.execute_reply.started":"2024-09-04T08:34:06.399150Z","shell.execute_reply":"2024-09-04T08:34:06.402244Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def gen_monets(photos):\n    photos = photos.to(config.device)\n    monets = model.photo_to_monet(photos)\n    return [TF.to_pil_image(inverse_transform(monet)) for monet in monets]","metadata":{"execution":{"iopub.status.busy":"2024-09-04T08:34:06.404352Z","iopub.execute_input":"2024-09-04T08:34:06.404624Z","iopub.status.idle":"2024-09-04T08:34:06.415423Z","shell.execute_reply.started":"2024-09-04T08:34:06.404600Z","shell.execute_reply":"2024-09-04T08:34:06.414592Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Save generated monets","metadata":{}},{"cell_type":"code","source":"!mkdir /tmp/images","metadata":{"execution":{"iopub.status.busy":"2024-09-04T08:34:06.416528Z","iopub.execute_input":"2024-09-04T08:34:06.417549Z","iopub.status.idle":"2024-09-04T08:34:07.437005Z","shell.execute_reply.started":"2024-09-04T08:34:06.417499Z","shell.execute_reply":"2024-09-04T08:34:07.435807Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"generated_monets = []\nfor photo_batch in tqdm(test_loader, \"generating monets\"):\n    monets = gen_monets(photo_batch)\n    generated_monets += monets\n\nfor idx, monet in enumerate(tqdm(generated_monets, \"saving images\")):\n    monet.save(f'/tmp/images/{idx}.jpg')","metadata":{"execution":{"iopub.status.busy":"2024-09-04T08:34:07.438813Z","iopub.execute_input":"2024-09-04T08:34:07.439696Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import shutil\nshutil.make_archive(\"/kaggle/working/images\", 'zip', \"/tmp/images\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}