{"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":"# Cycle GAN (pytorch lightning + Weights and Biases)    \n\n---  \n\n* *Comfort arises from familiarity or precisely knowing where's what!*\n\n* So, why not use the same framework for all your CV/AI projects with near perfect experiment logging while being lightning fast?\n\n\n> Aryan Garg   \n\n\n**Note:**\n1. You'll need an account on [Weights & Biases](https://wandb.ai/site). Don't worry, it's free! \n\n**References:**\n1. [Tutorial: Cycle GAN from Scratch by Song Seung Won](https://www.kaggle.com/code/songseungwon/cyclegan-tutorial-from-scratch-monet-to-photo)","metadata":{}},{"cell_type":"code","source":"SUBMIT_NB = True","metadata":{"execution":{"iopub.status.busy":"2023-05-11T04:05:44.665596Z","iopub.execute_input":"2023-05-11T04:05:44.666006Z","iopub.status.idle":"2023-05-11T04:05:44.681142Z","shell.execute_reply.started":"2023-05-11T04:05:44.665976Z","shell.execute_reply":"2023-05-11T04:05:44.680060Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Imports","metadata":{}},{"cell_type":"markdown","source":"### wandb (WARNING: interactive cell)","metadata":{}},{"cell_type":"code","source":"if not SUBMIT_NB:\n    !pip -qqq install wandb pytorch-lightning torchmetrics\n\n    import wandb\n    from pytorch_lightning.loggers import WandbLogger\n\n    wandb.login()","metadata":{"execution":{"iopub.status.busy":"2023-05-11T04:05:47.438951Z","iopub.execute_input":"2023-05-11T04:05:47.439301Z","iopub.status.idle":"2023-05-11T04:05:47.447227Z","shell.execute_reply.started":"2023-05-11T04:05:47.439271Z","shell.execute_reply":"2023-05-11T04:05:47.446184Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Lightning","metadata":{}},{"cell_type":"code","source":"try:\n    import lightning.pytorch as pl\nexcept:\n    print(\"[!] Couldn't find pytorch-lightning.\\nInstalling it...\\n\")\n    !pip install lightning\n    import lightning.pytorch as pl","metadata":{"execution":{"iopub.status.busy":"2023-05-11T04:05:47.835812Z","iopub.execute_input":"2023-05-11T04:05:47.836372Z","iopub.status.idle":"2023-05-11T04:06:18.026587Z","shell.execute_reply.started":"2023-05-11T04:05:47.836337Z","shell.execute_reply":"2023-05-11T04:06:18.025597Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from lightning.pytorch.utilities.model_summary import ModelSummary","metadata":{"execution":{"iopub.status.busy":"2023-05-11T04:06:18.028496Z","iopub.execute_input":"2023-05-11T04:06:18.028864Z","iopub.status.idle":"2023-05-11T04:06:18.036589Z","shell.execute_reply.started":"2023-05-11T04:06:18.028829Z","shell.execute_reply":"2023-05-11T04:06:18.033979Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from pytorch_lightning import seed_everything","metadata":{"execution":{"iopub.status.busy":"2023-05-11T04:06:18.037847Z","iopub.execute_input":"2023-05-11T04:06:18.038577Z","iopub.status.idle":"2023-05-11T04:06:19.242395Z","shell.execute_reply.started":"2023-05-11T04:06:18.038543Z","shell.execute_reply":"2023-05-11T04:06:19.241485Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Standard imports","metadata":{}},{"cell_type":"code","source":"import os\nimport shutil\nimport pathlib\n\nimport PIL\nfrom PIL import Image\nimport numpy as np\nimport cv2 as cv\nimport random\nimport matplotlib.pyplot as plt\nfrom tqdm.auto import tqdm\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch import optim\n\nfrom torch.utils.data import Dataset, DataLoader, random_split\n\nimport torchvision\nfrom torchvision import datasets","metadata":{"execution":{"iopub.status.busy":"2023-05-11T04:22:24.226417Z","iopub.execute_input":"2023-05-11T04:22:24.227144Z","iopub.status.idle":"2023-05-11T04:22:24.234110Z","shell.execute_reply.started":"2023-05-11T04:22:24.227110Z","shell.execute_reply":"2023-05-11T04:22:24.232778Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import torchvision.transforms as T\nfrom torchvision.transforms import Compose, ToTensor, Resize\nfrom torchvision.utils import make_grid","metadata":{"execution":{"iopub.status.busy":"2023-05-11T04:06:19.399385Z","iopub.execute_input":"2023-05-11T04:06:19.399784Z","iopub.status.idle":"2023-05-11T04:06:19.405600Z","shell.execute_reply.started":"2023-05-11T04:06:19.399748Z","shell.execute_reply":"2023-05-11T04:06:19.404620Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Torchmetrics \n\n*Not needed.*     \n\nBut just in case if we decide to use a metric to draw some insights!      \n\nEspecially: FID since the competition evaluates on **Memorization-Informed FID**","metadata":{}},{"cell_type":"code","source":"try:\n    import torchmetrics\nexcept:\n    print(f\"[!] Torchmetrics couldn't be imported.\\nInstalling...\")\n    !pip install torchmetrics","metadata":{"execution":{"iopub.status.busy":"2023-05-11T04:06:19.406912Z","iopub.execute_input":"2023-05-11T04:06:19.409125Z","iopub.status.idle":"2023-05-11T04:06:19.415396Z","shell.execute_reply.started":"2023-05-11T04:06:19.409097Z","shell.execute_reply":"2023-05-11T04:06:19.414507Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Custom Utilities (Not many lol)","metadata":{}},{"cell_type":"code","source":"# Folder Utilities ----------------------------\n\n## Create dir if it doesn't exist\ndef create_dir(dir_name):\n    if not os.path.exists(f'/content/{dir_name}'):\n        os.mkdir(f'/content/{dir_name}')\n\n## Delete dir: checkpoints\ndef delete_dir(dir_name):\n    if os.path.isdir(f'/content/{dir_name}'):\n        shutil.rmtree(f'/content/{dir_name}')","metadata":{"execution":{"iopub.status.busy":"2023-05-11T04:06:19.416955Z","iopub.execute_input":"2023-05-11T04:06:19.417641Z","iopub.status.idle":"2023-05-11T04:06:19.425392Z","shell.execute_reply.started":"2023-05-11T04:06:19.417609Z","shell.execute_reply":"2023-05-11T04:06:19.424499Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"--- \n\n## Config File (logged to wandb), Seeds & Devices","metadata":{}},{"cell_type":"code","source":"# Log this config file to wandb\nCONFIG = dict(\n    seed=42,\n    DATA_ROOT = '/kaggle/input/gan-getting-started/',\n    BATCH_SIZE = 32,\n    WORKERS = 2,\n    IMG_SIZE = (256,256,3),\n    NUM_EPOCHS = 20,\n    lr = 0.0002,\n    b1 = 0.5,\n    b2 = 0.999,\n    disc_steps=1,\n    checkpoint_path='/kaggle/working/',\n    )","metadata":{"execution":{"iopub.status.busy":"2023-05-11T04:06:19.427071Z","iopub.execute_input":"2023-05-11T04:06:19.427797Z","iopub.status.idle":"2023-05-11T04:06:19.435331Z","shell.execute_reply.started":"2023-05-11T04:06:19.427763Z","shell.execute_reply":"2023-05-11T04:06:19.434450Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"seed_everything(CONFIG['seed'])","metadata":{"execution":{"iopub.status.busy":"2023-05-11T04:06:19.436931Z","iopub.execute_input":"2023-05-11T04:06:19.437604Z","iopub.status.idle":"2023-05-11T04:06:19.452811Z","shell.execute_reply.started":"2023-05-11T04:06:19.437571Z","shell.execute_reply":"2023-05-11T04:06:19.451855Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# device = torch.device('cpu')\nif torch.cuda.is_available():\n    device = torch.device('cuda')\nelse:\n    device = torch.device('cpu')\n\ndevice","metadata":{"execution":{"iopub.status.busy":"2023-05-11T04:06:19.457847Z","iopub.execute_input":"2023-05-11T04:06:19.458149Z","iopub.status.idle":"2023-05-11T04:06:19.473371Z","shell.execute_reply.started":"2023-05-11T04:06:19.458125Z","shell.execute_reply":"2023-05-11T04:06:19.472244Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"---   \n\n## Transforms","metadata":{}},{"cell_type":"code","source":"train_transform_album = Compose([T.Resize((256,256)),ToTensor()])\n#         A.SmallestMaxSize(max_size=160),\n#         A.ShiftScaleRotate(shift_limit=0.05, scale_limit=0.05, rotate_limit=15, p=0.5),\n#         A.RandomCrop(height=128, width=128),\n#         A.RGBShift(r_shift_limit=15, g_shift_limit=15, b_shift_limit=15, p=0.5),\n#         A.RandomBrightnessContrast(p=0.5),\n#         A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),","metadata":{"execution":{"iopub.status.busy":"2023-05-11T04:06:19.475020Z","iopub.execute_input":"2023-05-11T04:06:19.475360Z","iopub.status.idle":"2023-05-11T04:06:19.480719Z","shell.execute_reply.started":"2023-05-11T04:06:19.475327Z","shell.execute_reply":"2023-05-11T04:06:19.479737Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"---  \n\n## Dataset   \n\n### !! Use LightningDataModule this time? !!\n\n\n**Reqs:**\n1. Sort images and load aligned dataset (domA_img, domB_img)","metadata":{}},{"cell_type":"code","source":"class MonetDataset(Dataset):\n    def __init__(self, \n                 root_dir: str = None, \n                 cust_transform = Compose([ToTensor()]), \n                 isTrain: bool = True):\n        \n        super(MonetDataset).__init__()\n        self.root = root_dir\n        self.files_Y = os.listdir(self.root+\"monet_jpg/\")\n        start = np.random.randint(0,6701)\n        end = start+300\n        self.files_X = os.listdir(self.root+\"photo_jpg/\")[start:end]\n        self.transform = cust_transform\n        self.isTrain = isTrain\n        \n    def __getitem__(self, idx):\n        photo_X = Image.open(self.root+\"photo_jpg/\"+self.files_X[idx])\n        photo_Y = Image.open(self.root+\"monet_jpg/\"+self.files_Y[idx])\n        \n        item_X = self.transform(photo_X)\n        item_Y = self.transform(photo_Y)\n        \n        return (item_X, item_Y) # Photo , Monet\n    \n    def __len__(self):\n#         print(\"Monet:\",len(self.files_Y),\"Photo:\",len(self.files_X))\n        return max(len(self.files_Y), len(self.files_X))","metadata":{"execution":{"iopub.status.busy":"2023-05-11T04:06:19.482402Z","iopub.execute_input":"2023-05-11T04:06:19.483182Z","iopub.status.idle":"2023-05-11T04:06:19.494135Z","shell.execute_reply.started":"2023-05-11T04:06:19.483147Z","shell.execute_reply":"2023-05-11T04:06:19.493289Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if not SUBMIT_NB:\n    all_data = MonetDataset(CONFIG['DATA_ROOT'], cust_transform=train_transform_album)\n    real, monet = next(iter(all_data))\n\n    image_grid = torch.cat((real, monet), 1)\n    plt.imshow(image_grid.cpu().permute(1,2,0))\n    plt.title(\"top: real | bottom: monet\")\n    plt.axis('off')\n    plt.show()","metadata":{"execution":{"iopub.status.busy":"2023-05-11T04:06:19.495827Z","iopub.execute_input":"2023-05-11T04:06:19.496192Z","iopub.status.idle":"2023-05-11T04:06:19.503555Z","shell.execute_reply.started":"2023-05-11T04:06:19.496161Z","shell.execute_reply":"2023-05-11T04:06:19.502532Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if not SUBMIT_NB:\n    print(real.shape, real.dtype)\n    print(len(all_data))","metadata":{"execution":{"iopub.status.busy":"2023-05-11T04:06:19.505547Z","iopub.execute_input":"2023-05-11T04:06:19.505874Z","iopub.status.idle":"2023-05-11T04:06:19.512561Z","shell.execute_reply.started":"2023-05-11T04:06:19.505850Z","shell.execute_reply":"2023-05-11T04:06:19.511513Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class MonetDataModule(pl.LightningDataModule):\n    def __init__(self, data_dir: str = None, batch_size: int = 1):\n        super().__init__()\n        self.data_dir = data_dir\n        self.batch_size = batch_size\n        \n    def prepare_data(self):\n        pass\n\n    def setup(self, stage: str):\n        if stage == 'fit':\n            self.train_data = MonetDataset(CONFIG['DATA_ROOT'], isTrain=True)\n        elif stage == 'test': \n            # Not implemented anything different from above (training dataset)... \n            # This is just to highlight how the datamodule can be used for general purposes.\n            self.test_data = MonetDataset(CONFIG['DATA_ROOT'], isTrain=False)\n            \n    def train_dataloader(self):\n        return DataLoader(self.train_data, shuffle=True, batch_size=self.batch_size)\n    \n    def test_dataloader(self):\n        return DataLoader(self.test_data, shuffle=True, batch_size=self.batch_size)\n\n    def teardown(self, stage: str):\n        pass","metadata":{"execution":{"iopub.status.busy":"2023-05-11T04:06:19.514260Z","iopub.execute_input":"2023-05-11T04:06:19.515027Z","iopub.status.idle":"2023-05-11T04:06:19.524664Z","shell.execute_reply.started":"2023-05-11T04:06:19.514991Z","shell.execute_reply":"2023-05-11T04:06:19.523864Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"---   \n\n## Architecture","metadata":{}},{"cell_type":"code","source":"!pip install torchsummary","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from torchvision import models\nfrom torchsummary import summary","metadata":{"execution":{"iopub.status.busy":"2023-05-11T04:28:13.782099Z","iopub.execute_input":"2023-05-11T04:28:13.782486Z","iopub.status.idle":"2023-05-11T04:28:25.312701Z","shell.execute_reply.started":"2023-05-11T04:28:13.782439Z","shell.execute_reply":"2023-05-11T04:28:25.311441Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Residual Generator","metadata":{}},{"cell_type":"markdown","source":"Note to self:    \n\nnn.conv2d(in_f, out_f, k,s,p)    \n\nk -> kernel_size    \ns -> stride (default: 1)    \np -> padding (default: 0)  ","metadata":{}},{"cell_type":"code","source":"class ResBlk(nn.Module):\n    def __init__(self, in_features):\n        super().__init__()\n        self.model = nn.Sequential(\n            # Alpha blk begin --------------\n            nn.ReflectionPad2d(1),\n            nn.Conv2d(in_features, in_features, 3, 1, 0),\n            nn.InstanceNorm2d(in_features),\n            # Alpha blk end   --------------\n            nn.ReLU(inplace=True),\n            nn.ReflectionPad2d(1),\n            nn.Conv2d(in_features, in_features, 3, 1, 0),\n            nn.InstanceNorm2d(in_features),\n        )\n    \n    def forward(self, x):\n        return self.model(x) + x","metadata":{"execution":{"iopub.status.busy":"2023-05-11T04:28:25.315587Z","iopub.execute_input":"2023-05-11T04:28:25.315992Z","iopub.status.idle":"2023-05-11T04:28:25.323956Z","shell.execute_reply.started":"2023-05-11T04:28:25.315953Z","shell.execute_reply":"2023-05-11T04:28:25.322979Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class ResGen(nn.Module):\n    def __init__(self, channels, out_features=64, num_residual_blocks=9):\n        super().__init__()\n        self.c = channels\n        self.out_f = out_features\n        self.num_resBlks = num_residual_blocks\n        \n        # Alpha blk\n        self.model = [nn.ReflectionPad2d(self.c), \n                      nn.Conv2d(self.c, self.out_f, 7, 1, 0), \n                      nn.InstanceNorm2d(self.out_f)]\n        self.model += [nn.ReLU(inplace=True)]\n        \n        # Downsampler: Add 2 Beta blocks\n        out_f2 = None\n        for i in range(2):\n            out_f2 = 2*self.out_f\n            self.model += [nn.Conv2d(self.out_f, out_f2, 3,2,1), \n                           nn.InstanceNorm2d(out_f2), \n                           nn.ReLU(inplace=True)]\n            self.out_f = out_f2\n            \n        # Add residual blocks defined in the cell above\n        for i in range(self.num_resBlks):\n            self.model += [ResBlk(in_features=out_f2)]\n            \n        # Upsampler: Add 2 Gamma Blocks\n        in_up = out_f2\n        for i in range(2):\n            out_f2 //= 2\n            self.model += [nn.Upsample(scale_factor=2),\n                           nn.Conv2d(in_up, out_f2, 3,1,1),\n                           nn.ReLU(inplace=True)\n                          ]\n            in_up = out_f2\n            \n        # Output layer:\n        self.model += [nn.ReflectionPad2d(self.c), \n                       nn.Conv2d(out_f2, self.c ,7, 1, 0), \n                       nn.Tanh()]\n        \n        self.model = nn.Sequential(*self.model)\n        \n    def forward(self, x):\n        return self.model(x)","metadata":{"execution":{"iopub.status.busy":"2023-05-11T04:28:25.325712Z","iopub.execute_input":"2023-05-11T04:28:25.326065Z","iopub.status.idle":"2023-05-11T04:28:25.340252Z","shell.execute_reply.started":"2023-05-11T04:28:25.326032Z","shell.execute_reply":"2023-05-11T04:28:25.339155Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"resgen = ResGen(channels=3).to(device)\nsummary(resgen, (3, 256, 256))","metadata":{"execution":{"iopub.status.busy":"2023-05-11T04:28:25.342883Z","iopub.execute_input":"2023-05-11T04:28:25.343356Z","iopub.status.idle":"2023-05-11T04:28:25.527679Z","shell.execute_reply.started":"2023-05-11T04:28:25.343324Z","shell.execute_reply":"2023-05-11T04:28:25.526663Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### PatchGAN Discriminator","metadata":{}},{"cell_type":"code","source":"class patchDisc(nn.Module):\n    def __init__(self, channels):\n        super().__init__()\n        \n        self.c = channels\n        def disc_block(in_f, out_f):\n            return [nn.Conv2d(in_f, out_f, 4,2,1), nn.InstanceNorm2d(out_f), nn.LeakyReLU(0.2, inplace=True)]\n        \n        # PatchGAN calculation done by hand\n        self.output_shape = (1, 16, 16)\n        \n        self.model = nn.Sequential(\n            *disc_block(self.c, 64),\n            *disc_block(64, 128),\n            *disc_block(128, 256),\n            *disc_block(256, 512),\n            nn.ZeroPad2d((1,0,1,0)),\n            nn.Conv2d(512, 1, 4,1,1)\n        )\n    \n        \n    def forward(self, x):\n        return self.model(x)","metadata":{"execution":{"iopub.status.busy":"2023-05-11T04:06:31.291937Z","iopub.execute_input":"2023-05-11T04:06:31.292755Z","iopub.status.idle":"2023-05-11T04:06:31.301909Z","shell.execute_reply.started":"2023-05-11T04:06:31.292713Z","shell.execute_reply":"2023-05-11T04:06:31.300933Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if not SUBMIT_NB:\n    pdisc = patchDisc(3).to(device)\n    summary(pdisc, (3,256,256))","metadata":{"execution":{"iopub.status.busy":"2023-05-11T04:06:31.303703Z","iopub.execute_input":"2023-05-11T04:06:31.304375Z","iopub.status.idle":"2023-05-11T04:06:31.313972Z","shell.execute_reply.started":"2023-05-11T04:06:31.304338Z","shell.execute_reply":"2023-05-11T04:06:31.312998Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"---   \n\n## Lightning Recipe","metadata":{}},{"cell_type":"code","source":"class LIT_cycle(pl.LightningModule):\n    \n    def __init__(self, Gen_XY=None, Gen_YX=None, D_X=None, D_Y=None, \n                 lr: float = 1e-3, b1: float = 0.5, b2: float = 0.999, channels=3, disc_steps=1):\n        super().__init__()\n        \n        assert Gen_XY is not None, \"Pass the generator: X -> Y !\"\n        assert Gen_YX is not None, \"Pass the generator: Y -> X !\"\n        assert D_X is not None, \"Pass the patch-discriminator: X !\"\n        assert D_Y is not None, \"Pass the patch-discriminator: Y !\"\n        \n        self.save_hyperparameters(ignore=[])\n        self.automatic_optimization = False\n        \n        self.g_xy = Gen_XY\n        self.g_yx = Gen_YX\n        self.dx = D_X\n        self.dy = D_Y\n        self.channels = channels\n        \n        self.criterion_GAN = torch.nn.MSELoss().to(device)\n        self.cycle_loss = torch.nn.L1Loss().to(device)\n        self.identity_loss = torch.nn.L1Loss().to(device)\n        \n        \n    def configure_optimizers(self):\n        lr = self.hparams.lr # .hparams is accessed from the args passed to this core Lightning module :)\n        b1 = self.hparams.b1\n        b2 = self.hparams.b2\n        \n        # chain the two generator optimizers so that:\n        # both generators are updated together (using the same optimizer instance)!\n        from itertools import chain\n        optim_g = torch.optim.Adam(chain(self.g_xy.parameters(), self.g_yx.parameters()), lr=lr, betas=(b1,b2))\n        \n        optim_dx = torch.optim.Adam(self.dx.parameters(), lr=lr, betas=(b1,b2))\n        optim_dy = torch.optim.Adam(self.dy.parameters(), lr=lr, betas=(b1,b2))\n        \n        # Second list is for returning any lr-schedulers that you might want to use!\n        return [optim_g, optim_dx, optim_dy], []\n    \n    \n    def forward(self, z, toDomain='x'):\n        if toDomain == 'x':\n            return self.g_yx(z)\n        else:\n            return self.g_xy(z)\n        \n\n    def training_step(self, batch, idx):\n        real_X, real_Y = batch\n        real_X, real_Y = real_X.to(device), real_Y.to(device)\n        \n        valid = torch.Tensor(np.ones((real_X.size(0), *self.dx.output_shape))).to(device)\n        fake = torch.Tensor(np.zeros((real_X.size(0), *self.dx.output_shape))).to(device)\n        \n        opt_g, opt_dx, opt_dy = self.optimizers()\n        \n        # Train G:\n        self.toggle_optimizer(opt_g)\n        \n        # TODO: Compute loss:\n        out_x = self.g_yx(real_X)\n        out_y = self.g_xy(real_Y)\n#         print(out_x.shape, out_y.shape)\n        \n        loss_id_X = self.identity_loss(out_x, real_X)\n        loss_id_Y = self.identity_loss(out_y, real_Y)\n        loss_identity = (loss_id_X + loss_id_Y)/2\n        \n        # GAN Loss\n        fake_Y = self.g_xy(real_X)\n        loss_GAN_XY = self.criterion_GAN(self.dy(fake_Y), valid) # tricking the 'fake-Y' into 'real-Y'\n        fake_X = self.g_yx(real_Y)\n        loss_GAN_YX = self.criterion_GAN(self.dx(fake_X), valid) # tricking the 'fake-X' into 'real-X'\n        \n        loss_GAN = (loss_GAN_XY + loss_GAN_YX)/2\n        \n        # Cycle Loss\n        recov_X = self.g_yx(fake_Y) # recov_X is fake-photo that is generated by fake-monet-drawing \n        loss_cycle_X = self.cycle_loss(recov_X, real_X) # Reduces the difference between the restored image and the real image\n        recov_Y = self.g_xy(fake_X)\n        loss_cycle_Y = self.cycle_loss(recov_Y, real_Y)\n        \n        loss_cycle = (loss_cycle_X + loss_cycle_Y)/2.\n        \n        # Total loss:         \n        g_loss = loss_GAN + (10.0*loss_cycle) + (5.0*loss_identity) # multiply by weights suggested by paper-authors\n        \n        # Gradient step:\n        self.manual_backward(g_loss)\n        opt_g.step()\n        opt_g.zero_grad()\n        \n        self.untoggle_optimizer(opt_g)\n        \n        # Train Discriminator\n        for i in range(self.hparams.disc_steps): # is kept at 1, usually\n            # Toggle dx\n            self.toggle_optimizer(opt_dx)\n            # Compute losses for dx:\n            loss_real = self.criterion_GAN(self.dx(real_X), valid) # Discriminate real images as real\n            loss_fake = self.criterion_GAN(self.dx(fake_X.detach()), fake) # Discriminate fake images as fake\n        \n            dx_loss = (loss_real + loss_fake)/2.\n            \n            self.manual_backward(dx_loss)\n            opt_dx.step()\n            opt_dx.zero_grad()\n            self.untoggle_optimizer(opt_dx)\n            \n            # Toggle dy\n            self.toggle_optimizer(opt_dy)\n            # Compute losses for dy\n            loss_real = self.criterion_GAN(self.dy(real_Y), valid) # Discriminate real images as real\n            loss_fake = self.criterion_GAN(self.dy(fake_Y.detach()), fake) # Discriminate fake images as fake\n        \n            dy_loss = (loss_real + loss_fake)/2.\n            \n            self.manual_backward(dy_loss)\n            opt_dy.step()\n            opt_dy.zero_grad()\n            self.untoggle_optimizer(opt_dy)\n            \n            tot_d_loss = (dx_loss + dy_loss)/2.\n        \n        if not SUBMIT_NB:\n            self.logger.experiment.log({\"Gen. Monet\":[wandb.Image(make_grid(fake_Y[0].cpu()), caption=\"Gen Monet\")]})\n            self.logger.experiment.log({\"Gen. Photo\":[wandb.Image(make_grid(fake_X[0].cpu()), caption=\"Gen Photo\")]})\n            self.log_dict({\"g_loss\": g_loss, \"dx_loss\": dx_loss, \"dy_loss\": dy_loss, \"tot_d_loss\": tot_d_loss}, \n                      on_step=True, \n                      on_epoch=True, \n                      prog_bar=True, \n                      logger=True)","metadata":{"execution":{"iopub.status.busy":"2023-05-11T04:28:46.672981Z","iopub.execute_input":"2023-05-11T04:28:46.673355Z","iopub.status.idle":"2023-05-11T04:28:46.699833Z","shell.execute_reply.started":"2023-05-11T04:28:46.673322Z","shell.execute_reply":"2023-05-11T04:28:46.698882Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Init Recipe\ncycle_gan = LIT_cycle(Gen_XY=ResGen(3), \n                      Gen_YX=ResGen(3), \n                      D_X=patchDisc(3), \n                      D_Y=patchDisc(3), \n                      lr=CONFIG['lr'], \n                      b1=CONFIG['b1'],\n                      b2=CONFIG['b2'],\n                      channels=CONFIG['IMG_SIZE'][2],\n                      disc_steps=CONFIG['disc_steps']\n                     )\n\n# summary = ModelSummary(cycle_gan, max_depth=-1)\n# print(summary)","metadata":{"execution":{"iopub.status.busy":"2023-05-11T04:28:48.161441Z","iopub.execute_input":"2023-05-11T04:28:48.162162Z","iopub.status.idle":"2023-05-11T04:28:48.557604Z","shell.execute_reply.started":"2023-05-11T04:28:48.162119Z","shell.execute_reply":"2023-05-11T04:28:48.556540Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"---   \n\n## Logger: Project Name, Run Name, Config file logging etc.","metadata":{}},{"cell_type":"code","source":"if not SUBMIT_NB:\n    wandb.login()","metadata":{"execution":{"iopub.status.busy":"2023-05-11T04:28:53.925160Z","iopub.execute_input":"2023-05-11T04:28:53.925543Z","iopub.status.idle":"2023-05-11T04:28:53.931492Z","shell.execute_reply.started":"2023-05-11T04:28:53.925510Z","shell.execute_reply":"2023-05-11T04:28:53.930389Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if not SUBMIT_NB:\n    wandb_logger = WandbLogger(project='M8-cycleGAN', \n                           name='exp-2_20eps',\n                           config=CONFIG,\n                           job_type='train',\n                           log_model=\"all\")","metadata":{"execution":{"iopub.status.busy":"2023-05-11T04:28:54.145246Z","iopub.execute_input":"2023-05-11T04:28:54.146031Z","iopub.status.idle":"2023-05-11T04:28:54.151651Z","shell.execute_reply.started":"2023-05-11T04:28:54.145990Z","shell.execute_reply":"2023-05-11T04:28:54.150497Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"---    \n\n## Trainer Callbacks","metadata":{}},{"cell_type":"code","source":"from pytorch_lightning import Callback\nfrom lightning.pytorch.callbacks import DeviceStatsMonitor, TQDMProgressBar, ModelCheckpoint, EarlyStopping, LearningRateMonitor\n\n# Checkpoint\ncheckpoint_callback = ModelCheckpoint(dirpath=CONFIG['checkpoint_path'],\n                                      filename='{epoch}-{g_loss:.3f}',\n                                      monitor='g_loss',\n                                      save_top_k=-1,\n                                      save_last=True,\n                                      save_weights_only=True,\n                                      verbose=True,\n                                      mode='min')\n\n# Exp2: Learning Rate Monitor\nlr_monitor = LearningRateMonitor(logging_interval='step', log_momentum=False)\n\n# Earlystopping\n# earlystopping = EarlyStopping(monitor='val_d_acc', patience=3, mode='min')","metadata":{"execution":{"iopub.status.busy":"2023-05-11T04:28:55.444318Z","iopub.execute_input":"2023-05-11T04:28:55.445332Z","iopub.status.idle":"2023-05-11T04:28:55.453044Z","shell.execute_reply.started":"2023-05-11T04:28:55.445286Z","shell.execute_reply":"2023-05-11T04:28:55.451872Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"---   \n\n## Trainer","metadata":{}},{"cell_type":"code","source":"if SUBMIT_NB:\n    logger_name = False\n    callbacks_lst = []\nelse:\n    logger_name = wandb_logger\n    callbacks_lst = [TQDMProgressBar(refresh_rate=100), checkpoint_callback, lr_monitor]\ntrainer = pl.Trainer(fast_dev_run=False,    # For debugging purposes\n                     log_every_n_steps=1,   # set the logging frequency\n                     accelerator='auto',    # Precedence: tpu > gpu >> cpu\n                     devices=\"auto\",        # all\n                     max_epochs= CONFIG['NUM_EPOCHS'], # CONFIG['NUM_EPOCHS'],\n                     callbacks=callbacks_lst,\n                     logger=logger_name,    # wandb <3 OR False for NB submission\n                    )","metadata":{"execution":{"iopub.status.busy":"2023-05-11T04:28:59.861840Z","iopub.execute_input":"2023-05-11T04:28:59.862226Z","iopub.status.idle":"2023-05-11T04:28:59.920972Z","shell.execute_reply.started":"2023-05-11T04:28:59.862194Z","shell.execute_reply":"2023-05-11T04:28:59.919994Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"---  \n\n## Training","metadata":{}},{"cell_type":"markdown","source":"**TODO:** Replace CustomDataset calling with: MonetDataModule Implementation","metadata":{}},{"cell_type":"code","source":"monet_data = MonetDataset(CONFIG['DATA_ROOT'], isTrain=True)\ntrain_loader = DataLoader(monet_data, batch_size=1, shuffle=True)","metadata":{"execution":{"iopub.status.busy":"2023-05-11T04:29:03.028569Z","iopub.execute_input":"2023-05-11T04:29:03.029197Z","iopub.status.idle":"2023-05-11T04:29:03.040143Z","shell.execute_reply.started":"2023-05-11T04:29:03.029160Z","shell.execute_reply":"2023-05-11T04:29:03.039146Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# trainer.fit(cycle_gan, train_loader)","metadata":{"execution":{"iopub.status.busy":"2023-05-11T04:29:04.031145Z","iopub.execute_input":"2023-05-11T04:29:04.032784Z","iopub.status.idle":"2023-05-11T04:29:04.037426Z","shell.execute_reply.started":"2023-05-11T04:29:04.032737Z","shell.execute_reply":"2023-05-11T04:29:04.036401Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Finish Logging","metadata":{}},{"cell_type":"code","source":"if not SUBMIT_NB:\n    wandb.finish()\n    print(\"Finished W&B session\")","metadata":{"execution":{"iopub.status.busy":"2023-05-11T04:29:05.997743Z","iopub.execute_input":"2023-05-11T04:29:05.998839Z","iopub.status.idle":"2023-05-11T04:29:06.003796Z","shell.execute_reply.started":"2023-05-11T04:29:05.998800Z","shell.execute_reply":"2023-05-11T04:29:06.002787Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"---\n\n\n## Submission (eval steps, logging & zipping results)","metadata":{}},{"cell_type":"code","source":"cycle_20 = LIT_cycle.load_from_checkpoint(\"/kaggle/input/cycle-gan-checkpoint-20-epochs/cycle_checkpoint_20.ckpt\", \n                                          map_location=device)\n# summary = ModelSummary(cycle_20, max_depth=-1)\n# print(summary)\n\n","metadata":{"execution":{"iopub.status.busy":"2023-05-11T04:29:12.320846Z","iopub.execute_input":"2023-05-11T04:29:12.321279Z","iopub.status.idle":"2023-05-11T04:29:12.576704Z","shell.execute_reply.started":"2023-05-11T04:29:12.321225Z","shell.execute_reply":"2023-05-11T04:29:12.575727Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from torchvision import models\nfrom torchsummary import summary\nsummary(cycle_20.g_xy, (3,256,256))","metadata":{"execution":{"iopub.status.busy":"2023-05-11T04:30:30.854779Z","iopub.execute_input":"2023-05-11T04:30:30.855158Z","iopub.status.idle":"2023-05-11T04:30:30.883602Z","shell.execute_reply.started":"2023-05-11T04:30:30.855127Z","shell.execute_reply":"2023-05-11T04:30:30.882466Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"! rm -rf ../images\n! rm -rf /kaggle/working/images.zip\n! mkdir ../images","metadata":{"execution":{"iopub.status.busy":"2023-05-11T04:39:01.666751Z","iopub.execute_input":"2023-05-11T04:39:01.667121Z","iopub.status.idle":"2023-05-11T04:39:04.730518Z","shell.execute_reply.started":"2023-05-11T04:39:01.667091Z","shell.execute_reply":"2023-05-11T04:39:04.729164Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"inf_transform = Compose([T.Resize((256,256)), ToTensor()])","metadata":{"execution":{"iopub.status.busy":"2023-05-11T04:38:06.200773Z","iopub.execute_input":"2023-05-11T04:38:06.201495Z","iopub.status.idle":"2023-05-11T04:38:06.207777Z","shell.execute_reply.started":"2023-05-11T04:38:06.201430Z","shell.execute_reply":"2023-05-11T04:38:06.206659Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"photo_dir = \"/kaggle/input/gan-getting-started/photo_jpg\"\ncycle_gan.eval()\nwith torch.inference_mode():\n    for i, img_file in enumerate(tqdm(os.listdir(photo_dir))):\n        img = inf_transform(Image.open(photo_dir+\"/\"+img_file)).to(device)\n        img = img[None, :, :, :]\n        prediction = cycle_20.g_xy(img).cpu()\n        prediction = prediction[0]\n        \n        im = T.ToPILImage()(prediction).convert('RGB')\n        im.save(\"../images/\" + str(i+1) + \".jpg\")   ","metadata":{"execution":{"iopub.status.busy":"2023-05-11T04:38:08.347077Z","iopub.execute_input":"2023-05-11T04:38:08.347455Z","iopub.status.idle":"2023-05-11T04:38:08.587513Z","shell.execute_reply.started":"2023-05-11T04:38:08.347423Z","shell.execute_reply":"2023-05-11T04:38:08.586543Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import shutil","metadata":{"execution":{"iopub.status.busy":"2023-05-11T04:38:10.704284Z","iopub.execute_input":"2023-05-11T04:38:10.704682Z","iopub.status.idle":"2023-05-11T04:38:10.709108Z","shell.execute_reply.started":"2023-05-11T04:38:10.704648Z","shell.execute_reply":"2023-05-11T04:38:10.708114Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"shutil.make_archive(\"/kaggle/working/images\", 'zip', \"/kaggle/images\")","metadata":{"execution":{"iopub.status.busy":"2023-05-11T04:38:11.179399Z","iopub.execute_input":"2023-05-11T04:38:11.179799Z","iopub.status.idle":"2023-05-11T04:38:11.196659Z","shell.execute_reply.started":"2023-05-11T04:38:11.179766Z","shell.execute_reply":"2023-05-11T04:38:11.195605Z"},"trusted":true},"execution_count":null,"outputs":[]}]}