{"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":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2023-06-02T17:29:43.975093Z","iopub.execute_input":"2023-06-02T17:29:43.975390Z","iopub.status.idle":"2023-06-02T17:29:49.936641Z","shell.execute_reply.started":"2023-06-02T17:29:43.975349Z","shell.execute_reply":"2023-06-02T17:29:49.935810Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip install pytorch_lightning==1.9.3;","metadata":{"execution":{"iopub.status.busy":"2023-06-02T17:29:55.412330Z","iopub.execute_input":"2023-06-02T17:29:55.412714Z","iopub.status.idle":"2023-06-02T17:30:11.519095Z","shell.execute_reply.started":"2023-06-02T17:29:55.412685Z","shell.execute_reply":"2023-06-02T17:30:11.517883Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import glob\nimport os\nimport shutil\n\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 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\ndevice = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\nprint(device)","metadata":{"execution":{"iopub.status.busy":"2023-06-02T17:30:11.522625Z","iopub.execute_input":"2023-06-02T17:30:11.522997Z","iopub.status.idle":"2023-06-02T17:30:37.609689Z","shell.execute_reply.started":"2023-06-02T17:30:11.522967Z","shell.execute_reply":"2023-06-02T17:30:37.608818Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CustomImageDataset(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)\n","metadata":{"execution":{"iopub.status.busy":"2023-06-02T17:30:37.611264Z","iopub.execute_input":"2023-06-02T17:30:37.611680Z","iopub.status.idle":"2023-06-02T17:30:37.618520Z","shell.execute_reply.started":"2023-06-02T17:30:37.611646Z","shell.execute_reply":"2023-06-02T17:30:37.616952Z"},"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        \n        # ensure images outside of training dataset are also of the same size\n        self.transform = T.Resize((target_dim, target_dim), antialias=True)\n        \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":{"execution":{"iopub.status.busy":"2023-06-02T17:30:37.622148Z","iopub.execute_input":"2023-06-02T17:30:37.622580Z","iopub.status.idle":"2023-06-02T17:30:37.635931Z","shell.execute_reply.started":"2023-06-02T17:30:37.622469Z","shell.execute_reply":"2023-06-02T17:30:37.634963Z"},"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 = CustomImageDataset(self.monet_filenames, self.transform, stage)\n            self.train_photo = CustomImageDataset(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 = CustomImageDataset(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        return loaders\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":{"execution":{"iopub.status.busy":"2023-06-02T17:30:37.637413Z","iopub.execute_input":"2023-06-02T17:30:37.637788Z","iopub.status.idle":"2023-06-02T17:30:37.654434Z","shell.execute_reply.started":"2023-06-02T17:30:37.637734Z","shell.execute_reply":"2023-06-02T17:30:37.653602Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"DM_CONFIG = {    \n    # Directories of the Monet Paintings and Photos\n    \"monet_dir\": os.path.join(\"./Dataset/gan-getting-started/monet_jpg\", \"*.jpg\"),\n    \"photo_dir\": os.path.join(\"./Dataset/gan-getting-started/photo_jpg\", \"*.jpg\"),\n    \n    \"loader_config\": {\n        # the number of subprocesses (excluding the main process) used for data loading\n        \"num_workers\": os.cpu_count(),\n        \n        # enable faster data transfer to GPU during training\n        \"pin_memory\": torch.cuda.is_available(),\n    },\n    \n    # the validation/test batch size (mainly for visualization purposes)\n    \"sample_size\": 5,\n    \n    # the training/prediction batch size\n    \"batch_size\": 1\n}","metadata":{"execution":{"iopub.status.busy":"2023-06-02T17:30:37.655622Z","iopub.execute_input":"2023-06-02T17:30:37.655942Z","iopub.status.idle":"2023-06-02T17:30:37.673873Z","shell.execute_reply.started":"2023-06-02T17:30:37.655913Z","shell.execute_reply":"2023-06-02T17:30:37.673030Z"},"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=(18, 8))\n    plt.imshow(img_grid)\n    plt.axis(\"off\")\n    plt.title(title)\n    plt.show()\n","metadata":{"execution":{"iopub.status.busy":"2023-06-02T17:30:37.675162Z","iopub.execute_input":"2023-06-02T17:30:37.676136Z","iopub.status.idle":"2023-06-02T17:30:37.692300Z","shell.execute_reply.started":"2023-06-02T17:30:37.676103Z","shell.execute_reply":"2023-06-02T17:30:37.691531Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def init_weights(m):\n    if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d, nn.InstanceNorm2d)):\n      torch.nn.init.normal_(m.weight, mean=0.0, std=0.02)\n      if m.bias is not None:\n        nn.init.constant_(m.bias, 0.0)","metadata":{"execution":{"iopub.status.busy":"2023-06-02T17:30:37.693348Z","iopub.execute_input":"2023-06-02T17:30:37.693617Z","iopub.status.idle":"2023-06-02T17:30:37.704190Z","shell.execute_reply.started":"2023-06-02T17:30:37.693596Z","shell.execute_reply":"2023-06-02T17:30:37.703296Z"},"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                # Apply Init Weights\n        #self.block.apply(init_weights)\n        \n    def forward(self, x):\n        return self.block(x)\n","metadata":{"execution":{"iopub.status.busy":"2023-06-02T17:30:37.705808Z","iopub.execute_input":"2023-06-02T17:30:37.706154Z","iopub.status.idle":"2023-06-02T17:30:37.720719Z","shell.execute_reply.started":"2023-06-02T17:30:37.706125Z","shell.execute_reply":"2023-06-02T17:30:37.719647Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def test_Downsampling():\n  # Test going from 3 in_channels -> 64 hidden_channels -> 128 out channels\n  # with batch size of 5 and image size 3 x 256 x 256\n  in_channels = 3\n  hidden_channels = 64\n  out_channels = 128\n  batch_size = 5\n  H = 256\n  W = 256\n  # Create Zero Tensor of desired input shape\n  x = torch.zeros(((batch_size, in_channels, H, W)))\n  print(\"Initial Size:\")\n  print(x.shape)\n  print()\n\n  # Downsample 1: define beginning in->out channels\n  model1 = Downsampling(in_channels, hidden_channels)\n  scores = model1(x)\n  # You should see size [batch_size, hidden_channels, H/2, W/2]\n  # Which = [5, 64, 128, 128] in this case\n  print(\"Downsample 1:\")\n  print(scores.size())\n  print()\n\n  # Downsampling 2, need to redefine input/output size\n  model2 = Downsampling(hidden_channels, out_channels)\n  scores2 = model2(scores)\n  # You should see size [batch_size, out_channels, H/4, W/4]\n  # Which = [5, 128, 64, 64] in this case\n  print(\"Downsample 2:\")\n  print(scores2.size())\n\n\ntest_Downsampling()","metadata":{"execution":{"iopub.status.busy":"2023-06-02T17:30:37.725741Z","iopub.execute_input":"2023-06-02T17:30:37.726511Z","iopub.status.idle":"2023-06-02T17:30:38.090408Z","shell.execute_reply.started":"2023-06-02T17:30:37.726478Z","shell.execute_reply":"2023-06-02T17:30:38.089386Z"},"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                # Apply Init Weights\n        #self.block.apply(init_weights)\n        \n    def forward(self, x):\n        return self.block(x)\n","metadata":{"execution":{"iopub.status.busy":"2023-06-02T17:30:38.092230Z","iopub.execute_input":"2023-06-02T17:30:38.092860Z","iopub.status.idle":"2023-06-02T17:30:38.101496Z","shell.execute_reply.started":"2023-06-02T17:30:38.092824Z","shell.execute_reply":"2023-06-02T17:30:38.100521Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def test_Upsampling():\n  # Test going from 128 in_channels -> 64 hidden_channels -> 3 out channels\n  # with batch size of 5 and input size 128 x 64 x 64\n  in_channels = 128\n  hidden_channels = 64\n  out_channels = 3\n  batch_size = 5\n  H = 64\n  W = 64\n  # Create Zero Tensor of desired input shape\n  x = torch.zeros(((batch_size, in_channels, H, W)))\n  print(\"Initial Size:\")\n  print(x.shape)\n  print()\n\n  # Upsample 1: define beginning in->out channels\n  model1 = Upsampling(in_channels, hidden_channels)\n  scores = model1(x)\n  # You should see size [batch_size, hidden_channels, H*2, W*2]\n  # Which = [5, 64, 128, 128] in this case\n  print(\"Upsample 1:\")\n  print(scores.size())\n  print()\n\n  # Upsampling 2, need to redefine input/output size\n  model2 = Upsampling(hidden_channels, out_channels)\n  scores2 = model2(scores)\n  # You should see size [batch_size, out_channels, H*4, W*4]\n  # Which = [5, 3, 256, 256] in this case\n  print(\"Upsample 2:\")\n  print(scores2.size())\n\n\ntest_Upsampling()","metadata":{"execution":{"iopub.status.busy":"2023-06-02T17:30:38.103075Z","iopub.execute_input":"2023-06-02T17:30:38.103566Z","iopub.status.idle":"2023-06-02T17:30:38.306341Z","shell.execute_reply.started":"2023-06-02T17:30:38.103533Z","shell.execute_reply":"2023-06-02T17:30:38.305185Z"},"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":{"execution":{"iopub.status.busy":"2023-06-02T17:30:38.308171Z","iopub.execute_input":"2023-06-02T17:30:38.308708Z","iopub.status.idle":"2023-06-02T17:30:38.316278Z","shell.execute_reply.started":"2023-06-02T17:30:38.308672Z","shell.execute_reply":"2023-06-02T17:30:38.315110Z"},"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            # downsampling path\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            \n            # residual blocks\n            *[ResBlock(hid_channels*4) for _ in range(num_resblocks)], # 256x64x64\n            \n            # upsampling path\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.\")\n","metadata":{"execution":{"iopub.status.busy":"2023-06-02T17:30:38.317842Z","iopub.execute_input":"2023-06-02T17:30:38.318177Z","iopub.status.idle":"2023-06-02T17:30:38.339339Z","shell.execute_reply.started":"2023-06-02T17:30:38.318146Z","shell.execute_reply":"2023-06-02T17:30:38.338561Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def test_UNet():\n  # Test with input image of size 3 x 256 x 256 with batch_size = 5\n  in_channels = 3\n  hidden_channels = 64\n  out_channels = 3\n  batch_size = 5\n  H = 256\n  W = 256\n  # Create Zero Tensor of batch x image size\n  image = torch.zeros(((batch_size, in_channels, H, W)))\n  print(\"Initial Size:\")\n  print(image.shape)\n  print()\n\n  # Apply UNet Generator!\n  model = UNetGenerator(hidden_channels, in_channels, out_channels)\n  generated_image = model(image)\n  # You should see that UNet ends with the same size image generated\n  # In this Case [batch_size, out_channels, H, W]\n  print(\"UNet Generation Size:\")\n  print(generated_image.size())\n\ntest_UNet()","metadata":{"execution":{"iopub.status.busy":"2023-06-02T17:30:38.342658Z","iopub.execute_input":"2023-06-02T17:30:38.342925Z","iopub.status.idle":"2023-06-02T17:30:40.260096Z","shell.execute_reply.started":"2023-06-02T17:30:38.342901Z","shell.execute_reply":"2023-06-02T17:30:40.257514Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def test_ResNet():\n  # Test with input image of size 3 x 256 x 256 with batch_size = 5\n  in_channels = 3\n  hidden_channels = 64\n  out_channels = 3\n  batch_size = 5\n  num_resblocks = 4\n  H = 256\n  W = 256\n  # Create Zero Tensor of batch x image size\n  image = torch.zeros(((batch_size, in_channels, H, W)))\n  print(\"Initial Size:\")\n  print(image.shape)\n  print()\n\n  # Apply UNet Generator!\n  model = ResNetGenerator(hidden_channels, in_channels, out_channels, num_resblocks)\n  generated_image = model(image)\n  # You should see that UNet ends with the same size image generated\n  # In this Case [batch_size, out_channels, H, W]\n  print(\"ResNet Generation Size:\")\n  print(generated_image.size())\n\ntest_ResNet()","metadata":{"execution":{"iopub.status.busy":"2023-06-02T17:30:40.261887Z","iopub.execute_input":"2023-06-02T17:30:40.262398Z","iopub.status.idle":"2023-06-02T17:30:44.569659Z","shell.execute_reply.started":"2023-06-02T17:30:40.262346Z","shell.execute_reply":"2023-06-02T17:30:44.568631Z"},"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":{"execution":{"iopub.status.busy":"2023-06-02T17:30:44.571235Z","iopub.execute_input":"2023-06-02T17:30:44.571873Z","iopub.status.idle":"2023-06-02T17:30:44.579733Z","shell.execute_reply.started":"2023-06-02T17:30:44.571837Z","shell.execute_reply":"2023-06-02T17:30:44.578781Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def test_Patch():\n  # Test with input image of size 3 x 256 x 256 with batch_size = 5\n  in_channels = 3\n  hidden_channels = 64\n  out_channels = 3\n  batch_size = 5\n  H = 256\n  W = 256\n  # Create Zero Tensor of batch x image size\n  image = torch.ones(((batch_size, in_channels, H, W)))\n  print(\"Initial Size:\")\n  print(image.shape)\n  print()\n\n  # Apply Patch Discriminator!\n  model = Discriminator(hidden_channels, in_channels)\n  probability_scores = model(image)\n  # You should see that Discriminator ends with \n  # probability matrix of size  [batch_size, 1, 30, 30]\n  print(\"PatchGan Probability Scores Size:\")\n  print(probability_scores.size())\n  print()\n\n  print(\"PatchGan Probability Scores:\")\n  print(probability_scores)\n\ntest_Patch()","metadata":{"execution":{"iopub.status.busy":"2023-06-02T17:30:44.581115Z","iopub.execute_input":"2023-06-02T17:30:44.581522Z","iopub.status.idle":"2023-06-02T17:30:45.036184Z","shell.execute_reply.started":"2023-06-02T17:30:44.581469Z","shell.execute_reply":"2023-06-02T17:30:45.035257Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# adapted from https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/master/util/image_pool.py\nclass 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)\n","metadata":{"execution":{"iopub.status.busy":"2023-06-02T17:30:45.037740Z","iopub.execute_input":"2023-06-02T17:30:45.038102Z","iopub.status.idle":"2023-06-02T17:30:45.047158Z","shell.execute_reply.started":"2023-06-02T17:30:45.038068Z","shell.execute_reply":"2023-06-02T17:30:45.046197Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"DM_CONFIG = {    \n    # Directories of the Monet Paintings and Photos\n    \"monet_dir\": os.path.join(\"/kaggle/input/gan-getting-started/monet_jpg\", \"*.jpg\"),\n\n    \"photo_dir\": os.path.join(\"/kaggle/input/gan-getting-started/photo_jpg\", \"*.jpg\"),\n                              \n    \n    \"loader_config\": {\n        # the number of subprocesses (excluding the main process) used for data loading\n        \"num_workers\": os.cpu_count(),\n        \n        # enable faster data transfer to GPU during training\n        \"pin_memory\": torch.cuda.is_available(),\n    },\n    \n    # the validation/test batch size (mainly for visualization purposes)\n    \"sample_size\": 5,\n    \n    # the training/prediction batch size\n    \"batch_size\": 1\n}","metadata":{"execution":{"iopub.status.busy":"2023-06-02T17:30:45.049298Z","iopub.execute_input":"2023-06-02T17:30:45.049872Z","iopub.status.idle":"2023-06-02T17:30:45.075698Z","shell.execute_reply.started":"2023-06-02T17:30:45.049840Z","shell.execute_reply":"2023-06-02T17:30:45.074858Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"dm_sample = CustomDataModule(**{k: v for k, v in DM_CONFIG.items()})\ndm_sample.setup(\"fit\")\n\ntrain_loader = dm_sample.train_dataloader()\nimgs = next(iter(train_loader))\n\nmonet_ds = train_loader[\"monet\"]\nphoto_ds = train_loader[\"photo\"]\n\n\n# Lengths of the Dataset\nprint(\"Number of Monet Paintings: \" + str(len(monet_ds)))\nprint(\"Number of Photos: \" + str(len(photo_ds)))\nprint()\n\n# Image size 3 x 256 x 256\nprint(\"Image Sizes: \")\nprint(next(iter(monet_ds))[0].shape)\nprint(next(iter(photo_ds))[0].shape)\nprint()\n\n# Store a set of example images to understand and visualize\nexample_monets = torch.empty((5, 3, 256, 256))\nexample_photos = torch.empty((5, 3, 256, 256))\n\n# Example of how to keep getting images form each Dataset\n# Note Batch Size = 1 so we get a batch of 1 on each next()\nfor i in range(5):\n  example_monets[i] = next(iter(monet_ds))\n  example_photos[i] = next(iter(photo_ds))\n\nprint(\"Example Batch Sizes:\")\nprint(example_monets.shape)\nprint(example_photos.shape)\n#show_img(imgs[\"monet\"], nrow=5, title=\"Augmented Monet Paintings\")\n#show_img(imgs[\"photo\"], nrow=5, title=\"Augmented Photos\")","metadata":{"execution":{"iopub.status.busy":"2023-06-02T17:30:45.076917Z","iopub.execute_input":"2023-06-02T17:30:45.077890Z","iopub.status.idle":"2023-06-02T17:30:56.756413Z","shell.execute_reply.started":"2023-06-02T17:30:45.077859Z","shell.execute_reply":"2023-06-02T17:30:56.755103Z"},"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=(18, 8))\n    plt.imshow(img_grid)\n    plt.axis(\"off\")\n    plt.title(title)\n    plt.show()\n\nshow_img(example_monets, nrow=5, title=\"Monet Paintings\")\nshow_img(example_photos, nrow=5, title=\"Photos\")","metadata":{"execution":{"iopub.status.busy":"2023-06-02T17:30:56.758078Z","iopub.execute_input":"2023-06-02T17:30:56.759034Z","iopub.status.idle":"2023-06-02T17:30:57.767110Z","shell.execute_reply.started":"2023-06-02T17:30:56.758982Z","shell.execute_reply":"2023-06-02T17:30:57.766258Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def init_weights(m):\n    if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d, nn.InstanceNorm2d)):\n      torch.nn.init.normal_(m.weight, mean=0.0, std=0.02)\n      if m.bias is not None:\n        nn.init.constant_(m.bias, 0.0)","metadata":{"execution":{"iopub.status.busy":"2023-06-02T17:30:57.768481Z","iopub.execute_input":"2023-06-02T17:30:57.769476Z","iopub.status.idle":"2023-06-02T17:30:57.776132Z","shell.execute_reply.started":"2023-06-02T17:30:57.769441Z","shell.execute_reply":"2023-06-02T17:30:57.774938Z"},"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                # Apply Init Weights\n        #self.block.apply(init_weights)\n        \n    def forward(self, x):\n        return self.block(x)\n","metadata":{"execution":{"iopub.status.busy":"2023-06-02T17:30:57.777795Z","iopub.execute_input":"2023-06-02T17:30:57.778382Z","iopub.status.idle":"2023-06-02T17:30:57.791615Z","shell.execute_reply.started":"2023-06-02T17:30:57.778333Z","shell.execute_reply":"2023-06-02T17:30:57.790374Z"},"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                # Apply Init Weights\n        #self.block.apply(init_weights)\n        \n    def forward(self, x):\n        return self.block(x)\n","metadata":{"execution":{"iopub.status.busy":"2023-06-02T17:30:57.793051Z","iopub.execute_input":"2023-06-02T17:30:57.793750Z","iopub.status.idle":"2023-06-02T17:30:57.804052Z","shell.execute_reply.started":"2023-06-02T17:30:57.793716Z","shell.execute_reply":"2023-06-02T17:30:57.803121Z"},"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":{"execution":{"iopub.status.busy":"2023-06-02T17:30:57.805573Z","iopub.execute_input":"2023-06-02T17:30:57.806242Z","iopub.status.idle":"2023-06-02T17:30:57.823440Z","shell.execute_reply.started":"2023-06-02T17:30:57.806209Z","shell.execute_reply":"2023-06-02T17:30:57.822605Z"},"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            # downsampling path\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            \n            # residual blocks\n            *[ResBlock(hid_channels*4) for _ in range(num_resblocks)], # 256x64x64\n            \n            # upsampling path\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.\")\n","metadata":{"execution":{"iopub.status.busy":"2023-06-02T17:30:57.825291Z","iopub.execute_input":"2023-06-02T17:30:57.825585Z","iopub.status.idle":"2023-06-02T17:30:57.847138Z","shell.execute_reply.started":"2023-06-02T17:30:57.825561Z","shell.execute_reply":"2023-06-02T17:30:57.846289Z"},"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":{"execution":{"iopub.status.busy":"2023-06-02T17:30:57.852477Z","iopub.execute_input":"2023-06-02T17:30:57.852789Z","iopub.status.idle":"2023-06-02T17:30:57.866970Z","shell.execute_reply.started":"2023-06-02T17:30:57.852765Z","shell.execute_reply":"2023-06-02T17:30:57.866158Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# adapted from https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/master/util/image_pool.py\nclass 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)\n","metadata":{"execution":{"iopub.status.busy":"2023-06-02T17:30:57.868273Z","iopub.execute_input":"2023-06-02T17:30:57.868686Z","iopub.status.idle":"2023-06-02T17:30:57.885654Z","shell.execute_reply.started":"2023-06-02T17:30:57.868656Z","shell.execute_reply":"2023-06-02T17:30:57.884868Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Base Dimensions to use\nhid_channels = 64\nin_channels = 3\nout_channels = 3\nnum_resblocks = 6\n\n# Generator to transform Photos to Monet Paintings\n#monet_generator = UNetGenerator(hid_channels, in_channels, out_channels)\n# Generator to transform Monet Paintings to Photos\n#photo_generator = UNetGenerator(hid_channels, in_channels, out_channels)\n\n\n\n# Generator to transform Photos to Monet Paintings\nmonet_generator = ResNetGenerator(hid_channels, in_channels, out_channels, num_resblocks)\n# Generator to transform Monet Paintings to Photos\nphoto_generator = ResNetGenerator(hid_channels, in_channels, out_channels, num_resblocks)\n\n# Give a batch of example photos to each generator\nto_monet = monet_generator(example_photos)\nto_photos = photo_generator(example_monets)\n\nshow_img(example_photos, nrow=5, title=\"Original Photos\")\nshow_img(to_monet, nrow=5, title=\"Generated Monets\")\nshow_img(example_monets, nrow=5, title=\"Original Monets\")\nshow_img(to_photos, nrow=5, title=\"Generated Photos\")","metadata":{"execution":{"iopub.status.busy":"2023-06-02T17:30:57.886995Z","iopub.execute_input":"2023-06-02T17:30:57.887748Z","iopub.status.idle":"2023-06-02T17:31:11.165134Z","shell.execute_reply.started":"2023-06-02T17:30:57.887716Z","shell.execute_reply":"2023-06-02T17:31:11.164218Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Discriminator to identify real Monet Paintings\nmonet_discriminator = Discriminator(hid_channels)\n# Discriminator to identify real photos\nphoto_discriminator = Discriminator(hid_channels)\n\n# Give a real photo to each discriminator (first one photo shown above)\nreal_monet_scores = monet_discriminator(example_monets[0])\nreal_photo_scores = photo_discriminator(example_photos[0])\n\n# Sizes are randomly initialized\nprint(\"Monet Discriminator Size and Scores\")\nprint(real_monet_scores.shape)\nprint(real_monet_scores)\nprint()\nprint(\"Photo Discriminator Size and Scores\")\nprint(real_photo_scores.shape)\nprint(real_photo_scores)","metadata":{"execution":{"iopub.status.busy":"2023-06-02T17:31:11.166674Z","iopub.execute_input":"2023-06-02T17:31:11.167266Z","iopub.status.idle":"2023-06-02T17:31:11.397479Z","shell.execute_reply.started":"2023-06-02T17:31:11.167232Z","shell.execute_reply":"2023-06-02T17:31:11.396485Z"},"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\")\n\ntrain_loader = dm_sample.train_dataloader()\nimgs = next(iter(train_loader))\n\n#show_img(imgs[\"monet\"], nrow=5, title=\"Augmented Monet Paintings\")\n#show_img(imgs[\"photo\"], nrow=5, title=\"Augmented Photos\")","metadata":{"execution":{"iopub.status.busy":"2023-06-02T17:31:11.401037Z","iopub.execute_input":"2023-06-02T17:31:11.401685Z","iopub.status.idle":"2023-06-02T17:31:11.438089Z","shell.execute_reply.started":"2023-06-02T17:31:11.401652Z","shell.execute_reply":"2023-06-02T17:31:11.437195Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# How good are we at tricking the discriminator\ndef get_adv_loss(fake_image_discriminator_score):\n    # A perfect fake will have the discriminator produce all 1s\n    real_labels = torch.ones_like(fake_image_discriminator_score)\n    # Retrun mse between actual scores and 1s\n    return F.mse_loss(fake_image_discriminator_score, real_labels)","metadata":{"execution":{"iopub.status.busy":"2023-06-02T17:31:11.439582Z","iopub.execute_input":"2023-06-02T17:31:11.439926Z","iopub.status.idle":"2023-06-02T17:31:11.445118Z","shell.execute_reply.started":"2023-06-02T17:31:11.439895Z","shell.execute_reply":"2023-06-02T17:31:11.443896Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# How similar are our twice transformed photos\n# i.e real -> generated -> recon should have similar real and recon\ndef get_cycle_loss(real, recon, lambda_cycle):\n    # Simply return reg * L1 loss between recon image and real image\n    return lambda_cycle * F.l1_loss(recon, real)","metadata":{"execution":{"iopub.status.busy":"2023-06-02T17:31:11.446936Z","iopub.execute_input":"2023-06-02T17:31:11.447313Z","iopub.status.idle":"2023-06-02T17:31:11.463408Z","shell.execute_reply.started":"2023-06-02T17:31:11.447279Z","shell.execute_reply":"2023-06-02T17:31:11.462561Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# How similar the real and identity images are\n# Given a photo in the correct domain, our generator should do very little\ndef get_idt_loss(real, identity_image, lambda_cycle, lambda_idt):\n    # Simply return reg * L1 loss between real and identity\n    return lambda_idt * lambda_cycle * F.l1_loss(identity_image, real)","metadata":{"execution":{"iopub.status.busy":"2023-06-02T17:31:11.465269Z","iopub.execute_input":"2023-06-02T17:31:11.466055Z","iopub.status.idle":"2023-06-02T17:31:11.477602Z","shell.execute_reply.started":"2023-06-02T17:31:11.466021Z","shell.execute_reply":"2023-06-02T17:31:11.476757Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Computes the Loss of BOTH Generators, given\n# - Original Photos from dataset, real_monet and real_photo\n# - Generated Fakes, fake_monet and fake_photo\n# - Reconstructed Photos, reconstructed_monet, reconstructed_photo\n# - Identities, monet_identity and photo_identity\n# - Discrim Scores of fake images, fake_monet_discrim_score, fake_photo_discrim_score\n# - Model Hyperparams lambda_cycle and lambda_idt\ndef get_gen_loss(real_monet, real_photo, fake_monet, fake_photo, reconstructed_monet, reconstructed_photo,\n                 monet_identity, photo_identity, fake_monet_discrim_score, fake_photo_discrim_score, lambda_cycle, lambda_idt):\n  \n\n    # STEP 1: Compute Adversarial Loss using fakes and discriminator scores\n    # Conceptually - how good the generatros are at tricking the discriminator\n    # Adv Loss of monet generator\n    adv_loss_photo_to_monet = get_adv_loss(fake_monet_discrim_score)\n    # Adv Loss of photo generator\n    adv_loss_monet_to_photo = get_adv_loss(fake_photo_discrim_score)\n    total_adv_loss = adv_loss_photo_to_monet + adv_loss_monet_to_photo\n            \n    # STEP 2: Compute Cycle Loss using\n    # Cycle loss of the original monet compared to reconstructed monet\n    cycle_loss_MPM = get_cycle_loss(real_monet, reconstructed_monet, lambda_cycle[0])\n    # Cycle loss of the original photo compared to reconstructed photo\n    cycle_loss_PMP = get_cycle_loss(real_photo, reconstructed_photo, lambda_cycle[1])\n    total_cycle_loss = cycle_loss_MPM + cycle_loss_PMP\n\n    # STEP 3: Compute Identity Loss using reals and identities\n    # Identity Loss of the real monet and monet identity\n    idt_loss_MM = get_idt_loss(real_monet, monet_identity, lambda_cycle[0], lambda_idt)\n    # Identity Loss of the real photo and identity photo\n    idt_loss_PP = get_idt_loss(real_photo, photo_identity, lambda_cycle[1], lambda_idt)\n    total_idt_loss = idt_loss_MM + idt_loss_PP\n            \n    # FINAL STEP: Sum losses to get total generator loss\n    gen_loss = total_adv_loss + total_idt_loss + total_cycle_loss\n    return gen_loss","metadata":{"execution":{"iopub.status.busy":"2023-06-02T17:31:11.479189Z","iopub.execute_input":"2023-06-02T17:31:11.479840Z","iopub.status.idle":"2023-06-02T17:31:11.491031Z","shell.execute_reply.started":"2023-06-02T17:31:11.479808Z","shell.execute_reply":"2023-06-02T17:31:11.490169Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Computes loss of Single Discriminator given fake and real discriminator scores\ndef get_disc_loss(real_discrim_score, fake_discrim_score):\n    # STEP 1: Calculate real discrim loss\n    # Our discriminator should output all 1s as this is a real image from the dataset\n    real_labels = torch.ones_like(real_discrim_score)\n    # Return mse between actual scores and 1s\n    real_loss = F.mse_loss(real_discrim_score, real_labels)\n        \n    # STEP 2: Calculate fake discrim loss\n    # Our discriminator should output all 0s as this is a generated image from the dataset\n    fake_labels = torch.zeros_like(fake_discrim_score)\n    # Return mse between actual scores and 1s\n    fake_loss = F.mse_loss(fake_discrim_score, fake_labels)\n        \n    # Return average of losses\n    disc_loss = (fake_loss + real_loss) * 0.5\n    return disc_loss","metadata":{"execution":{"iopub.status.busy":"2023-06-02T17:31:11.493483Z","iopub.execute_input":"2023-06-02T17:31:11.494229Z","iopub.status.idle":"2023-06-02T17:31:11.511407Z","shell.execute_reply.started":"2023-06-02T17:31:11.494196Z","shell.execute_reply":"2023-06-02T17:31:11.510588Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CycleGAN_n(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        #self.optimizer_disc = torch.optim.Adam(self.discriminator.parameters(), lr=0.001)\n\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                #nn.init.xavier_normal_(m.weight)\n\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        \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, optimizer_idx=0)\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        # ...\n\n        # train discriminators\n        self.toggle_optimizer(opt_disc, optimizer_idx=1)\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\n        #self.toggle_optimizer(optimizer_idx)\n\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}.\")\n\n\n    def toggle_optimizer(self, optimizer, optimizer_idx):\n        for i, param_group in enumerate(optimizer.param_groups):\n            if i == optimizer_idx:\n                param_group[\"lr\"] = self.lr_schedulers()[i].get_last_lr()[0]\n            else:\n                param_group[\"lr\"] = 0.0\n","metadata":{"execution":{"iopub.status.busy":"2023-06-02T17:31:11.512900Z","iopub.execute_input":"2023-06-02T17:31:11.513286Z","iopub.status.idle":"2023-06-02T17:31:11.554255Z","shell.execute_reply.started":"2023-06-02T17:31:11.513251Z","shell.execute_reply":"2023-06-02T17:31:11.553234Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CycleGAN(L.LightningModule):\n    def __init__(\n        self,\n        in_channels,\n        out_channels,\n        hid_channels,\n        num_resblocks,\n        optimizer,\n        lr,\n        betas,\n        lambda_idt,\n        lambda_cycle,\n        num_epochs,\n        monet_gen = None,\n        photo_gen = None,\n        monet_disc = None,\n        photo_disc = None\n    ):\n        super().__init__()\n        self.save_hyperparameters(ignore=[\"optimizer\"])\n        self.optimizer = optimizer\n        # Necessary to set as we are using multiple optimizers\n        self.automatic_optimization = False\n\n        # Step for Graphs\n        self.step_num = 0\n        \n        # INITIALIZE OUR GENERATORS AND DISCRIMINATORS!!!!\n        # Weights will automatically be initialized here\n\n        # Generators:\n        # Generator to transform Photos to Monet Paintings\n        if monet_gen == None:\n          self.monet_generator = ResNetGenerator(hid_channels, in_channels, out_channels, num_resblocks)\n        else:\n          self.monet_generator = monet_gen\n  \n        # Generator to transform Monet Paintings into Photos\n        if photo_gen == None:\n          self.photo_generator = ResNetGenerator(hid_channels, in_channels, out_channels, num_resblocks)\n        else:\n          self.photo_generator = photo_gen\n\n        # Discriminators\n        # Discriminator to identify real Monet Paintings\n        if monet_disc == None:\n          self.monet_discriminator = Discriminator(hid_channels)\n        else:\n          self.monet_discriminator = monet_disc\n\n        # Discriminator to identify real Photos\n        if photo_disc == None:\n          self.photo_discriminator = Discriminator(hid_channels)\n        else:\n          self.photo_discriminator = photo_disc\n\n        # Metrics for Loss graphs\n        self.losses = {\n            \"gen_loss\": [],\n            \"disc_loss_M\": [],\n            \"disc_loss_P\": [],\n        }\n        \n    # As our predominant focus is creating the best Monet Painting Generator\n    # Our forward simply passes in an image to the Monet Generator\n    def forward(self, img):\n        return self.monet_generator(img)   \n    \n\n    # Configure the Optimizers for both our generator and discriminator\n    def configure_optimizers(self):\n        # Load hyperparams\n        opt_config = {\n            \"lr\": self.hparams.lr,\n            \"betas\": self.hparams.betas,\n        }\n        # Set up optimizer for generators\n        opt_gen = self.optimizer(\n            list(self.monet_generator.parameters()) + list(self.photo_generator.parameters()),\n            **opt_config,\n        )\n        # Set up optimizer for discriminators\n        opt_disc = self.optimizer(\n            list(self.monet_discriminator.parameters()) + list(self.photo_discriminator.parameters()),\n            **opt_config,\n        )\n        optimizers = [opt_gen, opt_disc]\n        # No schedulers for now\n        schedulers = []\n        return optimizers, schedulers\n    \n    def training_step(self, batch, batch_idx):\n        # Begin by storing batches of real monets and photos from the dataset\n        real_monet = batch[\"monet\"]\n        real_photo = batch[\"photo\"]\n\n        # Store the optimizers for the generators and discriminators\n        gen_optimizer, disc_optimizer = self.optimizers()\n\n        # STEP 1: Use Generators to produce fake images in both domains\n        # Fake monet generated from the real photo\n        fake_monet = self.monet_generator(real_photo)\n        # Fake photo generated from the real monet\n        fake_photo = self.photo_generator(real_monet)\n         \n        # STEP 2: Use generators and fakes to reconstruct original inputs\n        # Reconstruction of original monet painting\n        reconstructed_monet = self.monet_generator(fake_photo)\n        # Reconstruction of original photo\n        reconstructed_photo = self.photo_generator(fake_monet)\n\n        # STEP 3: Construct identities by feeding real inputs into the generators\n        # Monet identity: feed real monet into monet generator\n        monet_identity = self.monet_generator(real_monet)\n        # Photo identity: feed real photo into photo generator\n        photo_identity = self.photo_generator(real_photo)\n\n        # STEP 4: Compute Discriminator Scores to see how well our\n        # Discriminators differentiate real and fake images\n        # Scores of how good the fake generated monets are\n        fake_monet_discrim_score = self.monet_discriminator(fake_monet)\n        # Scores of how good the fake generated photos are\n        fake_photo_discrim_score = self.photo_discriminator(fake_photo)\n        # Scores of how good the real monets are\n        real_monet_discrim_score = self.monet_discriminator(real_monet)\n        # Scores of how good the real photos are\n        real_photo_discrim_score = self.photo_discriminator(real_photo)\n    \n        # STEP 5: Use Generator Loss to train Generators\n        # Standard pipeline: Loss -> zero_grad -> backward -> step\n        self.toggle_optimizer(gen_optimizer, 0)\n        gen_loss = get_gen_loss(real_monet, real_photo, fake_monet, fake_photo, reconstructed_monet, reconstructed_photo,\n                 monet_identity, photo_identity, fake_monet_discrim_score, fake_photo_discrim_score, self.hparams.lambda_cycle, self.hparams.lambda_idt)      \n        gen_optimizer.zero_grad()\n        self.manual_backward(gen_loss)\n        gen_optimizer.step()\n        self.untoggle_optimizer(0)\n        \n        # STEP 5: Use Discriminator Loss to train Discriminators\n        # Standard pipeline: Losses -> zero_grad -> backward -> step\n        self.toggle_optimizer(disc_optimizer, 1)\n        # Need to detatch otherwise backward is happening twice\n        monet_disc_loss = get_disc_loss(real_monet_discrim_score, self.monet_discriminator(fake_monet.detach()))\n        photo_disc_loss = get_disc_loss(real_photo_discrim_score, self.photo_discriminator(fake_photo.detach()))\n        disc_optimizer.zero_grad()\n        self.manual_backward(monet_disc_loss)\n        self.manual_backward(photo_disc_loss)\n        disc_optimizer.step()\n        self.untoggle_optimizer(1)\n        \n        # Done with training step\n        # Log metrics and losses\n        metrics = {\n            \"gen_loss\": gen_loss,\n            \"disc_loss_M\": monet_disc_loss,\n            \"disc_loss_P\": photo_disc_loss,\n        }\n        self.log_dict(metrics, on_step=False, on_epoch=True, prog_bar=True)\n\n        # Add to Losses 4 times each epoch\n        if (self.step % 1757 == 0):\n          self.losses[\"gen_loss\"].append(float(gen_loss))\n          self.losses[\"disc_loss_M\"].append(float(monet_disc_loss))\n          self.losses[\"disc_loss_P\"].append(float(photo_disc_loss))\n\n        self.step += 1\n        \n        \n    # Output current progress of generator on a validation set\n    def validation_step(self, batch, batch_idx):\n        self.display_results(batch, batch_idx, \"validate\")\n    \n    # Output current progress of generator on a test set\n    def test_step(self, batch, batch_idx):\n        self.display_results(batch, batch_idx, \"test\")\n        \n    # Predict the results of a batch\n    def predict_step(self, batch, batch_idx):\n        return self(batch)\n    \n    # Function to Visualize Results\n    def display_results(self, batch, batch_idx, stage):\n        # Store batch of real input photos and corresponding fake generated monets\n        real_photos = batch\n        fake_monets = self(real_photos)\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        # Display current image results\n        show_img(\n            torch.cat([real_photos, fake_monets], dim=0),\n            nrow=len(real_photos),\n            title=title,\n        )\n    \n    def on_train_epoch_start(self):\n        # Update Progress Bar\n        self.log(\"lr\", self.hparams.lr, on_step=False, on_epoch=True, prog_bar=True)\n\n        # Reset Step\n        self.step = 0\n        \n    def on_train_epoch_end(self):\n        # Print Current State of Models\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        \n        print(\"Training ended.\")\n        \n    def on_predict_epoch_end(self, results):\n        predictions = self.trainer.predict_loop.predictions\n        print(f\"Done predicting\")","metadata":{"execution":{"iopub.status.busy":"2023-06-02T17:31:11.555751Z","iopub.execute_input":"2023-06-02T17:31:11.556134Z","iopub.status.idle":"2023-06-02T17:31:11.584850Z","shell.execute_reply.started":"2023-06-02T17:31:11.556104Z","shell.execute_reply":"2023-06-02T17:31:11.584060Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"\nMODEL_CONFIG = {\n    # the type of generator, and the number of residual blocks if ResNet generator is used\n    #\"gen_name\": \"resnet\", # types: 'unet', 'resnet'\n    # In channels and out channels\n    \"in_channels\": 3,\n    \"out_channels\": 3,\n    \"num_resblocks\": 6,\n    \n    # the number of filters in the first layer for the generators and discriminators\n    \"hid_channels\": 64,\n    \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\n    # using Adam\n    \"optimizer\": torch.optim.Adam,\n    \n    # the learning rate and beta parameters for the Adam optimizer\n    \"lr\": 2e-4,\n    \"betas\": (0.5, 0.999),\n    \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    \n    # the size of the buffer that stores previously generated images\n    #\"buffer_size\": 100,\n    \n    # the number of epochs for training\n    \"num_epochs\": 5,\n    \n    # the number of epochs before starting the learning rate decay\n    #\"decay_epochs\": 18,\n}","metadata":{"execution":{"iopub.status.busy":"2023-06-02T17:31:11.586004Z","iopub.execute_input":"2023-06-02T17:31:11.586480Z","iopub.status.idle":"2023-06-02T17:31:11.604418Z","shell.execute_reply.started":"2023-06-02T17:31:11.586440Z","shell.execute_reply":"2023-06-02T17:31:11.603612Z"},"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    \"precision\": 16,\n\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,\n    \"limit_predict_batches\": 1.0,\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\": 2,\n}","metadata":{"execution":{"iopub.status.busy":"2023-06-02T17:31:11.609689Z","iopub.execute_input":"2023-06-02T17:31:11.610214Z","iopub.status.idle":"2023-06-02T17:31:11.623996Z","shell.execute_reply.started":"2023-06-02T17:31:11.610177Z","shell.execute_reply":"2023-06-02T17:31:11.623199Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"dm = CustomDataModule(**DM_CONFIG)\nmodel = CycleGAN(**MODEL_CONFIG)\ntrainer = L.Trainer(**TRAIN_CONFIG)\n\ntrainer.fit(model, datamodule=dm)","metadata":{"execution":{"iopub.status.busy":"2023-06-02T17:31:11.625289Z","iopub.execute_input":"2023-06-02T17:31:11.625674Z","iopub.status.idle":"2023-06-02T17:54:35.073469Z","shell.execute_reply.started":"2023-06-02T17:31:11.625642Z","shell.execute_reply":"2023-06-02T17:54:35.072412Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Load trainer with loaded models\ndm_predict = CustomDataModule(**DM_CONFIG)\nmodel_predict = CycleGAN(**MODEL_CONFIG, monet_gen=monet_generator, photo_gen=photo_generator, monet_disc=monet_discriminator, photo_disc=photo_discriminator)\ntrainer_predict = L.Trainer(**TRAIN_CONFIG)\n\n# Generate all monet paintings from training photos\npredictions = trainer_predict.predict(model_predict, datamodule=dm_predict)","metadata":{"execution":{"iopub.status.busy":"2023-06-02T17:54:35.077630Z","iopub.execute_input":"2023-06-02T17:54:35.077952Z","iopub.status.idle":"2023-06-02T17:55:54.895746Z","shell.execute_reply.started":"2023-06-02T17:54:35.077917Z","shell.execute_reply":"2023-06-02T17:55:54.894697Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"os.makedirs(\"../images\", exist_ok=True)\n\nidx = 0\nfor img in predictions:\n    # Extract 3x256x256 image\n    pred_img = img[0]\n    # Save image to desired directory\n    save_image(\n            pred_img.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\")\nprint(\"Submission Complete\")","metadata":{"execution":{"iopub.status.busy":"2023-06-02T17:59:29.763139Z","iopub.execute_input":"2023-06-02T17:59:29.763530Z","iopub.status.idle":"2023-06-02T17:59:56.053964Z","shell.execute_reply.started":"2023-06-02T17:59:29.763499Z","shell.execute_reply":"2023-06-02T17:59:56.052196Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}