{"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":"# **Photos To Monet**","metadata":{"execution":{"iopub.status.busy":"2023-10-29T01:53:44.433717Z","iopub.execute_input":"2023-10-29T01:53:44.433954Z","iopub.status.idle":"2023-10-29T01:53:44.438541Z","shell.execute_reply.started":"2023-10-29T01:53:44.433931Z","shell.execute_reply":"2023-10-29T01:53:44.437480Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip install deepspeed","metadata":{"execution":{"iopub.status.busy":"2023-10-29T01:53:44.440776Z","iopub.execute_input":"2023-10-29T01:53:44.441073Z","iopub.status.idle":"2023-10-29T01:54:11.063704Z","shell.execute_reply.started":"2023-10-29T01:53:44.441048Z","shell.execute_reply":"2023-10-29T01:54:11.062721Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import glob\nimport os\nimport shutil\nimport deepspeed as ds\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport pytorch_lightning as L\nimport torch\nimport torch.nn.functional as F\nimport torchvision.transforms as T\nfrom pytorch_lightning.utilities import CombinedLoader\nfrom torch import nn\nfrom torch.utils.data import DataLoader, Dataset\nfrom torchvision.io import read_image\nfrom torchvision.utils import make_grid, save_image\n_ = L.seed_everything(0, workers=True)","metadata":{"execution":{"iopub.status.busy":"2023-10-29T01:54:11.065944Z","iopub.execute_input":"2023-10-29T01:54:11.066314Z","iopub.status.idle":"2023-10-29T01:54:30.088703Z","shell.execute_reply.started":"2023-10-29T01:54:11.066278Z","shell.execute_reply":"2023-10-29T01:54:30.087932Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Data Preprocessing ","metadata":{}},{"cell_type":"code","source":"def show_img(img_tensor, nrow, title=\"\"):\n    img_tensor = img_tensor.detach().cpu() * 0.5 + 0.5\n    img_grid = make_grid(img_tensor, nrow=nrow).permute(1, 2, 0)\n    plt.figure(figsize=(10, 7))\n    plt.imshow(img_grid)\n    plt.axis(\"off\")\n    plt.title(title)\n    plt.show()","metadata":{"execution":{"iopub.status.busy":"2023-10-29T01:54:30.089823Z","iopub.execute_input":"2023-10-29T01:54:30.090113Z","iopub.status.idle":"2023-10-29T01:54:30.095718Z","shell.execute_reply.started":"2023-10-29T01:54:30.090081Z","shell.execute_reply":"2023-10-29T01:54:30.094867Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Augmenting the images.**\n\nBefore loading the datasets, we define CustomTransform for image augmentation. This improves learning by introducing more variety in the images during training instead of learning from the same set of images, especially when we only have 300 Monet paintings.","metadata":{}},{"cell_type":"code","source":"class CustomTransform(object):\n    def __init__(self, load_dim=286, target_dim=256):\n        self.transform_train = T.Compose([\n            T.Resize((load_dim, load_dim), antialias=True),\n            T.RandomCrop((target_dim, target_dim)),\n            T.RandomHorizontalFlip(p=0.5),\n            T.ColorJitter(brightness=0.2, contrast=0.2,\n                          saturation=0.2, hue=0.1),\n        ])\n        self.transform = T.Resize((target_dim, target_dim), antialias=True)   \n    def __call__(self, img, stage):\n        if stage == \"fit\":\n            img = self.transform_train(img)\n        else:\n            img = self.transform(img)\n        return img * 2 - 1","metadata":{"execution":{"iopub.status.busy":"2023-10-29T01:54:30.097988Z","iopub.execute_input":"2023-10-29T01:54:30.098250Z","iopub.status.idle":"2023-10-29T01:54:30.109799Z","shell.execute_reply.started":"2023-10-29T01:54:30.098226Z","shell.execute_reply":"2023-10-29T01:54:30.109061Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Storing dataset**","metadata":{}},{"cell_type":"code","source":"class CustomDataset(Dataset):\n    def __init__(self, filenames, transform, stage):\n        self.filenames = filenames\n        self.transform = transform\n        self.stage = stage\n        \n    def __len__(self):\n        return len(self.filenames)\n    \n    def __getitem__(self, idx):\n        img_name = self.filenames[idx]\n        img = read_image(img_name) / 255.0\n        return self.transform(img, stage=self.stage)","metadata":{"execution":{"iopub.status.busy":"2023-10-29T01:54:30.110848Z","iopub.execute_input":"2023-10-29T01:54:30.111071Z","iopub.status.idle":"2023-10-29T01:54:30.122324Z","shell.execute_reply.started":"2023-10-29T01:54:30.111050Z","shell.execute_reply":"2023-10-29T01:54:30.121581Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Iterating through the datasets**\n\nTo prepare the datasets, we load them into DataLoader separately, which can then iterate through the datasets as needed. Because the training dataset contains both the Monet paintings and photos, we pass both dataloaders into CombinedLoader for training. ","metadata":{}},{"cell_type":"code","source":"DEBUG = False\n\nDM_CONFIG = {    \n    \"monet_dir\": os.path.join(\"/kaggle/input/gan-getting-started/monet_jpg\", \"*.jpg\"),\n    \"photo_dir\": os.path.join(\"/kaggle/input/gan-getting-started/photo_jpg\", \"*.jpg\"),\n    \n    \"loader_config\": {\n        \"num_workers\": os.cpu_count(),\n        \"pin_memory\": torch.cuda.is_available(),\n    },\n    \"sample_size\": 5,\n    \"batch_size\": 1 if not DEBUG else 1,\n}","metadata":{"execution":{"iopub.status.busy":"2023-10-29T01:54:30.123426Z","iopub.execute_input":"2023-10-29T01:54:30.124264Z","iopub.status.idle":"2023-10-29T01:54:30.146934Z","shell.execute_reply.started":"2023-10-29T01:54:30.124233Z","shell.execute_reply":"2023-10-29T01:54:30.146221Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CustomDataModule(L.LightningDataModule):\n    def __init__(\n        self,\n        monet_dir,\n        photo_dir, \n        loader_config,\n        sample_size,\n        batch_size,\n    ):\n        super().__init__()\n        self.loader_config = loader_config\n        self.sample_size = sample_size\n        self.batch_size = batch_size       \n            \n        # store file paths\n        self.monet_filenames = sorted(glob.glob(monet_dir))\n        self.photo_filenames = sorted(glob.glob(photo_dir))\n        \n        # define transformations for image augmentation\n        self.transform = CustomTransform()\n        \n    def setup(self, stage):\n        if stage == \"fit\":\n            self.train_monet = CustomDataset(self.monet_filenames, self.transform, stage)\n            self.train_photo = CustomDataset(self.photo_filenames, self.transform, stage)\n            \n        if stage in [\"fit\", \"test\", \"predict\"]:\n            # to be used for test and prediction datasets as well\n            self.valid_photo = CustomDataset(self.photo_filenames, self.transform, None)\n            \n    def train_dataloader(self):\n        loader_config = {\n            \"shuffle\": True,\n            \"drop_last\": True,\n            \"batch_size\": self.batch_size,\n            **self.loader_config,\n        }\n        loader_monet = DataLoader(self.train_monet, **loader_config)\n        loader_photo = DataLoader(self.train_photo, **loader_config)\n        loaders = {\"monet\": loader_monet, \"photo\": loader_photo}\n        return CombinedLoader(loaders, mode=\"max_size_cycle\")\n    \n    def val_dataloader(self):\n        return DataLoader(\n            self.valid_photo,\n            batch_size=self.sample_size,\n            **self.loader_config,\n        )\n    \n    def test_dataloader(self):\n        return self.val_dataloader()\n    \n    def predict_dataloader(self):\n        return DataLoader(\n            self.valid_photo,\n            batch_size=self.batch_size,\n            **self.loader_config,\n        )","metadata":{"execution":{"iopub.status.busy":"2023-10-29T01:54:30.148230Z","iopub.execute_input":"2023-10-29T01:54:30.148573Z","iopub.status.idle":"2023-10-29T01:54:30.168433Z","shell.execute_reply.started":"2023-10-29T01:54:30.148540Z","shell.execute_reply":"2023-10-29T01:54:30.167703Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"We check that the datamodule defined is working as intended by visualizing samples of the images below.","metadata":{}},{"cell_type":"code","source":"dm_sample = CustomDataModule(batch_size=5, **{k: v for k, v in DM_CONFIG.items() if k != \"batch_size\"})\ndm_sample.setup(\"fit\")\ntrain_loader = dm_sample.train_dataloader()\nimgs = next(iter(train_loader))\nshow_img(imgs[\"monet\"], nrow=5, title=\"Augmented Monet Paintings\")\nshow_img(imgs[\"photo\"], nrow=5, title=\"Augmented Photos\")","metadata":{"execution":{"iopub.status.busy":"2023-10-29T01:54:30.169418Z","iopub.execute_input":"2023-10-29T01:54:30.169776Z","iopub.status.idle":"2023-10-29T01:54:37.879584Z","shell.execute_reply.started":"2023-10-29T01:54:30.169743Z","shell.execute_reply":"2023-10-29T01:54:37.878605Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Building GAN Architecture","metadata":{}},{"cell_type":"markdown","source":"**U-Net Generator**","metadata":{}},{"cell_type":"markdown","source":"A common architecture for the CycleGAN generator is the U-Net. U-Net is a network which consists of a sequence of downsampling blocks followed by a sequence of upsampling blocks, giving it the U-shaped architecture. In the upsampling path, we concatenate the outputs of the upsampling blocks and the outputs of the downsampling blocks symmetrically. This can be seen as a kind of skip connection, facilitating information flow in deep networks and reducing the impact of vanishing gradients.","metadata":{}},{"cell_type":"markdown","source":"**ResNet generator**","metadata":{}},{"cell_type":"markdown","source":"Similar to the U-Net architecture, the ResNet generator consists of the downsampling path and upsampling path. The difference is that the ResNet generator does not have the long skip connections from the concatenations of outputs. Instead, the ResNet generator uses residual blocks between the two paths. These residual blocks have short skip connections where the original input is added to the output.","metadata":{}},{"cell_type":"markdown","source":"**Downsampling blocks**","metadata":{}},{"cell_type":"markdown","source":"The downsampling blocks use convolution layers to increase the number of feature maps while reducing the dimensions of the 2D image.","metadata":{}},{"cell_type":"code","source":"class Downsampling(nn.Module):\n    def __init__(\n        self,\n        in_channels, # 輸入通道數\n        out_channels, # 輸出通道數\n        kernel_size=4, # mask為4*4矩陣\n        stride=2, # mask在進行卷積時的移動步幅\n        padding=1, #mask在卷積時會讓圖片外框補像素(輸入輸出的大小會變一樣)\n        norm=True, #true則進行歸一化\n        lrelu=True,#true則使用激活函數LeakyReLU 否則使用ReLU\n    ):\n        super().__init__()\n        self.block = nn.Sequential( #名為nn.Sequential的容器 內有nn.Conv2d一個2D卷積層\n            nn.Conv2d(in_channels, out_channels,\n                      kernel_size=kernel_size, stride=stride, padding=padding, bias=not norm),\n        )\n        if norm:\n            self.block.append(nn.InstanceNorm2d(out_channels, affine=True))\n        if lrelu is not None:\n            self.block.append(nn.LeakyReLU(0.2, True) if lrelu else nn.ReLU(True))\n        \n    def forward(self, x): #輸入x給self.block返回self.block(x)\n        return self.block(x)","metadata":{"execution":{"iopub.status.busy":"2023-10-29T01:54:37.881050Z","iopub.execute_input":"2023-10-29T01:54:37.881350Z","iopub.status.idle":"2023-10-29T01:54:37.889427Z","shell.execute_reply.started":"2023-10-29T01:54:37.881324Z","shell.execute_reply":"2023-10-29T01:54:37.888494Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Upsampling blocks**\n\nOn the other hand, the upsampling blocks contain transposed convolution layers, which combine the learned features to output an image with the original size.","metadata":{}},{"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, #drop層:防止過度擬合(訓練中選擇一部份神經元輸出為零)\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: #如dropout為true\n            self.block.append(nn.Dropout(0.5)) #則添加dropout(以0.5的機率讓神經元輸出為零)\n        self.block.append(nn.ReLU(True)) #必定加上一個ReLU\n        \n    def forward(self, x):\n        return self.block(x)","metadata":{"execution":{"iopub.status.busy":"2023-10-29T01:54:37.893906Z","iopub.execute_input":"2023-10-29T01:54:37.894191Z","iopub.status.idle":"2023-10-29T01:54:37.902837Z","shell.execute_reply.started":"2023-10-29T01:54:37.894166Z","shell.execute_reply":"2023-10-29T01:54:37.902055Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Residual blocks**\n\nAs described above, the residual blocks have convolution layers where the original input is added to the output.","metadata":{}},{"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) #加上ResBlock輸出的值 減輕梯度消失的問題","metadata":{"execution":{"iopub.status.busy":"2023-10-29T01:54:37.903944Z","iopub.execute_input":"2023-10-29T01:54:37.904210Z","iopub.status.idle":"2023-10-29T01:54:37.917281Z","shell.execute_reply.started":"2023-10-29T01:54:37.904186Z","shell.execute_reply":"2023-10-29T01:54:37.916303Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Building the generator**\n\nWith the building blocks defined, we can now build our CycleGAN generator. For reference, the output size of each block is commented below.","metadata":{}},{"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 通道數x圖片大小\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        ) # 將圖片用下採樣變成1*1大小\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        ) #將1*1圖片變成目標大小\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(), #確保像素值在-1到1之間 輸出為3x256x256(三原色所以是3)\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]) #將skip倒轉並去掉最後一個元素 給予上採樣執行\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) #將上採樣特徵圖x和下採樣特徵圖skip進行融合 有助細節補充\n    \nclass ResNetGenerator(nn.Module):\n    def __init__(self, hid_channels, in_channels, out_channels, num_resblocks):\n        super().__init__()\n        self.model = nn.Sequential(\n            nn.ReflectionPad2d(3), #反射填充圖像外框\n            Downsampling(in_channels, hid_channels, #進行下採樣減小圖像並增加通道\n                         kernel_size=7, stride=1, padding=0, lrelu=False), # 64x256x256\n            Downsampling(hid_channels, hid_channels*2, kernel_size=3, lrelu=False), # 128x128x128\n            Downsampling(hid_channels*2, hid_channels*4, kernel_size=3, lrelu=False), # 256x64x64\n            *[ResBlock(hid_channels*4) for _ in range(num_resblocks)], # 256x64x64 建立多個ResBlock改善圖像品質\n            Upsampling(hid_channels*4, hid_channels*2, kernel_size=3, output_padding=1), # 128x128x128\n            Upsampling(hid_channels*2, hid_channels, kernel_size=3, output_padding=1), # 64x256x256\n            nn.ReflectionPad2d(3),\n            nn.Conv2d(hid_channels, out_channels, kernel_size=7, stride=1, padding=0), # 3x256x256 卷積轉換成最後輸出圖像\n            nn.Tanh(),\n        )\n        \n    def forward(self, x):\n        return self.model(x)\n    \ndef get_gen(gen_name, hid_channels, num_resblocks, in_channels=3, out_channels=3):\n    if gen_name == \"unet\":\n        return UNetGenerator(hid_channels, in_channels, out_channels)\n    elif gen_name == \"resnet\":\n        return ResNetGenerator(hid_channels, in_channels, out_channels, num_resblocks)\n    else:\n        raise NotImplementedError(f\"Generator name '{gen_name}' not recognized.\") #根據輸入選擇要用哪個生成器","metadata":{"execution":{"iopub.status.busy":"2023-10-29T01:54:37.918780Z","iopub.execute_input":"2023-10-29T01:54:37.919066Z","iopub.status.idle":"2023-10-29T01:54:37.936905Z","shell.execute_reply.started":"2023-10-29T01:54:37.919039Z","shell.execute_reply":"2023-10-29T01:54:37.935868Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Patch Gan Generator**\n\nUnlike conventional networks that output a single probability of the input image being real or fake, CycleGAN uses the PatchGAN discriminator that outputs a matrix of values. Intuitively, each value of the output matrix checks the corresponding portion of the input image. Values closer to 1 indicate real classification and values closer to 0 indicate fake classification.","metadata":{}},{"cell_type":"markdown","source":"**Building the discriminator**\n\nIn general, the PatchGAN discriminator consists of a sequence of convolution layers, which can be built using the downsampling blocks defined earlier.","metadata":{}},{"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-10-29T01:54:37.938205Z","iopub.execute_input":"2023-10-29T01:54:37.938497Z","iopub.status.idle":"2023-10-29T01:54:37.951703Z","shell.execute_reply.started":"2023-10-29T01:54:37.938471Z","shell.execute_reply":"2023-10-29T01:54:37.950836Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Image buffer**\n\nThe original CycleGAN implementation updates the discriminator using a history of generated images instead of the latest images generated. This is done by setting up an image buffer that stores previously generated images. With probability 0.5, each newly generated image is swapped with a previously generated image stored in the buffer. This stabilizes training by giving the discriminator access to past information.","metadata":{}},{"cell_type":"code","source":"class ImageBuffer(object): #保持訓練穩定性:可以儲存之前生成的圖像 有心圖像時有50%機率進行代替\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            # 緩衝區圖像還沒滿的話\n            if self.curr_cap < self.buffer_size:\n                self.curr_cap += 1\n                self.buffer.append(img)\n                return_imgs.append(img)\n            else:\n                p = np.random.uniform(low=0., high=1.)\n                \n                # swap images between input and buffer with probability 0.5\n                if p > 0.5:\n                    idx = np.random.randint(low=0, high=self.buffer_size)\n                    tmp = self.buffer[idx].clone()\n                    self.buffer[idx] = img\n                    return_imgs.append(tmp)\n                else:\n                    return_imgs.append(img)\n        return torch.cat(return_imgs, dim=0)","metadata":{"execution":{"iopub.status.busy":"2023-10-29T01:54:37.953138Z","iopub.execute_input":"2023-10-29T01:54:37.953441Z","iopub.status.idle":"2023-10-29T01:54:37.963988Z","shell.execute_reply.started":"2023-10-29T01:54:37.953406Z","shell.execute_reply":"2023-10-29T01:54:37.962913Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**CycleGan**\n\nWith the generator and discriminator architectures defined, we can now build CycleGAN, which consists of two generators and two discriminators:\n\n* Generator for photo-to-Monet translation (gen_PM).\n* Generator for Monet-to-photo translation (gen_MP).\n* Discriminator for Monet paintings (disc_M).\n* Discriminator for photos (disc_P).\nWe use the Adam optimizer for model training. To optimize the parameters, we need to define the loss functions:\n\n* **Discriminator loss**:- For real images fed into the discriminator, the output matrix is compared against a matrix of 1s using the mean squared error. For fake images, the output matrix is compared against a matrix of 0s. This suggests that to minimize loss, the perfect discriminator outputs a matrix of 1s for real images and a matrix of 0s for fake images.\n* **Generator loss:-** This is composed of three different loss functions below.\n     *          *Adversarial loss* :-Fake images are fed into the discriminator and the output matrix is compared against a matrix of 1s using the mean squared error. To minimize loss, the generator needs to 'fool' the discriminator into thinking that the fake images are real and output a matrix of 1s.\n    *         *Identity loss*:- When a Monet painting is fed into the photo-to-Monet generator, we should get back the same Monet painting because nothing needs to be transformed. The same applies for photos fed into the Monet-to-photo generator. To encourage identity mapping, the difference in pixel values between the input image and generated image is measured using the l1 loss.\n    *         *Cycle loss* When a Monet painting is fed into the Monet-to-photo generator, and the generated image is fed back into the photo-to-Monet generator, it should transform back into the original Monet painting. The same applies for photos passed to the two generators to get back the original photos. To preserve information throughout this cycle, the l1 loss is used to measure the difference between the original image and the reconstructed image.","metadata":{}},{"cell_type":"markdown","source":"**Building the CycleGAN model**","metadata":{}},{"cell_type":"code","source":"MODEL_CONFIG = { #模型配置設定\n    # the type of generator, and the number of residual blocks if ResNet generator is used\n    \"gen_name\": \"unet\", # types: 'unet', 'resnet'\n    \"num_resblocks\": 6,\n    # the number of filters in the first layer for the generators and discriminators\n    \"hid_channels\": 64,\n    # using DeepSpeed's FusedAdam (currently GPU only) is slightly faster\n    \"optimizer\": ds.ops.adam.FusedAdam if torch.cuda.is_available() else torch.optim.Adam,\n    # the learning rate and beta parameters for the Adam optimizer\n    \"lr\": 2e-4,\n    \"betas\": (0.5, 0.999),\n    # the weights used in the identity loss and cycle loss\n    \"lambda_idt\": 0.5,\n    \"lambda_cycle\": (10, 10), # (MPM direction, PMP direction)\n    # the size of the buffer that stores previously generated images\n    \"buffer_size\": 100,\n    # the number of epochs for training\n    \"num_epochs\": 18 if not DEBUG else 2,\n    # the number of epochs before starting the learning rate decay\n    \"decay_epochs\": 18 if not DEBUG else 1,\n}","metadata":{"execution":{"iopub.status.busy":"2023-10-29T01:54:37.965273Z","iopub.execute_input":"2023-10-29T01:54:37.965574Z","iopub.status.idle":"2023-10-29T01:54:37.978549Z","shell.execute_reply.started":"2023-10-29T01:54:37.965536Z","shell.execute_reply":"2023-10-29T01:54:37.977639Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CycleGAN(L.LightningModule):\n    def __init__(\n        self,\n        gen_name,\n        num_resblocks,\n        hid_channels,\n        optimizer,\n        lr,\n        betas,\n        lambda_idt,\n        lambda_cycle,\n        buffer_size,\n        num_epochs,\n        decay_epochs,\n    ):\n        super().__init__()\n        self.save_hyperparameters(ignore=[\"optimizer\"])\n        self.optimizer = optimizer\n        self.automatic_optimization = False\n        \n        # define generators and discriminators\n        self.gen_PM = get_gen(gen_name, hid_channels, num_resblocks)\n        self.gen_MP = get_gen(gen_name, hid_channels, num_resblocks)\n        self.disc_M = Discriminator(hid_channels)\n        self.disc_P = Discriminator(hid_channels)\n        \n        # initialize buffers to store fake images\n        self.buffer_fake_M = ImageBuffer(buffer_size)\n        self.buffer_fake_P = ImageBuffer(buffer_size)\n        \n    def forward(self, img):\n        return self.gen_PM(img)   #forward用輸入圖像生成新圖像\n            \n    def init_weights(self): #初始化模型權重\n        def init_fn(m):\n            if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d, nn.InstanceNorm2d)):\n                nn.init.normal_(m.weight, 0.0, 0.02)\n                if m.bias is not None:\n                    nn.init.constant_(m.bias, 0.0)\n        \n        for net in [self.gen_PM, self.gen_MP, self.disc_M, self.disc_P]:\n            net.apply(init_fn)\n        \n    def setup(self, stage):\n        if stage == \"fit\":\n            self.init_weights()\n            print(\"Model initialized.\") #初始化後顯示提示句子\n            \n    def get_lr_scheduler(self, optimizer): #學習速度調節器\n        def lr_lambda(epoch):\n            len_decay_phase = self.hparams.num_epochs - self.hparams.decay_epochs + 1.0\n            curr_decay_step = max(0, epoch - self.hparams.decay_epochs + 1.0)\n            val = 1.0 - curr_decay_step / len_decay_phase\n            return max(0.0, val)\n        \n        return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lr_lambda)\n    \n    def configure_optimizers(self): #優化器和學習率調度器\n        opt_config = {\n            \"lr\": self.hparams.lr,\n            \"betas\": self.hparams.betas,\n        }\n        opt_gen = self.optimizer(\n            list(self.gen_PM.parameters()) + list(self.gen_MP.parameters()),\n            **opt_config,\n        )\n        opt_disc = self.optimizer(\n            list(self.disc_M.parameters()) + list(self.disc_P.parameters()),\n            **opt_config,\n        )\n        optimizers = [opt_gen, opt_disc]\n        schedulers = [self.get_lr_scheduler(opt) for opt in optimizers]\n        return optimizers, schedulers\n        \n    def adv_criterion(self, y_hat, y): #使用均方誤差生成圖像與真實的差距\n        return F.mse_loss(y_hat, y)\n    \n    def recon_criterion(self, y_hat, y): #使用絕對值誤差生成圖像與真實的差距\n        return F.l1_loss(y_hat, y)\n    \n    def get_adv_loss(self, fake, disc): #fake為假的disc判別器\n        fake_hat = disc(fake) #判別器判別假圖的評分 0到1 越像真越大\n        real_labels = torch.ones_like(fake_hat) #把所有假圖標籤設為1\n        adv_loss = self.adv_criterion(fake_hat, real_labels) #計算上面兩個的對抗損失\n        return adv_loss #因為希望假圖越像真越好 判別器給出的機率最好為1 因此把所有假圖標籤設為1設為最佳狀態\n    \n    def get_idt_loss(self, real, idt, lambda_cycle): #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): #real為真實的recon為循環一圈的\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) #假莫內圖和判別器M的對抗損失\n        adv_loss_MP = self.get_adv_loss(self.fake_P, self.disc_P) #假風景照和判別器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) #真圖都標上1\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) #假圖都標上0\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) #真與假莫內圖計算判別器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) #真與假風景照計算判別器P的損失\n    \n    def training_step(self, batch, batch_idx): #訓練步驟\n        self.real_M = batch[\"monet\"] #獲取真實風景照與真實莫內圖\n        self.real_P = batch[\"photo\"] \n        opt_gen, opt_disc = self.optimizers() #獲取鑑別器與生成器的優化器\n\n        # generate fake images\n        self.fake_M = self.gen_PM(self.real_P) #生成器生成假莫內圖\n        self.fake_P = self.gen_MP(self.real_M) #生成器生成假風景照\n        \n        # generate identity images\n        self.idt_M = self.gen_PM(self.real_M) #莫內圖生成器將真莫內圖轉成假莫內圖的身份損失\n        self.idt_P = self.gen_MP(self.real_P) #風景照生成器將真風景照轉成假風景照的身份損失\n        \n        # reconstruct images\n        self.recon_M = self.gen_PM(self.fake_P) #莫內圖生成器將假風景照轉成假莫內圖的循環損失\n        self.recon_P = self.gen_MP(self.fake_M) #風景照生成器將假莫內圖轉成假風景照的循環損失\n    \n        # train generators\n        self.toggle_optimizer(opt_gen) #切換到生成器的優化器\n        gen_loss = self.get_gen_loss() #計算生成器的總損失(含對抗損失、循環一致性損失、身份損失，要最小化)\n        opt_gen.zero_grad() #梯度歸零\n        self.manual_backward(gen_loss) #根據損失 計算生成器梯度\n        opt_gen.step() #使用生成優化器更新生成器參數 根據梯度進行更新\n        self.untoggle_optimizer(opt_gen) #切換到判別器的優化器\n        \n        # train discriminators\n        self.toggle_optimizer(opt_disc) #切換到判別器的優化器 \n        disc_loss_M = self.get_disc_loss_M() #兩個判別器的損失\n        disc_loss_P = self.get_disc_loss_P()\n        opt_disc.zero_grad() #梯度清零\n        self.manual_backward(disc_loss_M) #下面和上面一樣\n        self.manual_backward(disc_loss_P)\n        opt_disc.step()\n        self.untoggle_optimizer(opt_disc)\n        \n        # 紀錄訓練時的損失\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        # 更新學習率\n        for sch in self.lr_schedulers():\n            sch.step()\n        \n        # 印出當前訓練周期的狀態\n        logged_values = self.trainer.progress_bar_metrics\n        print(\n            f\"Epoch {self.current_epoch+1}\",\n            *[f\"{k}: {v:.5f}\" for k, v in logged_values.items()],\n            sep=\" - \",\n        )\n        \n    def on_train_end(self):\n        print(\"Training ended.\") #印出一條消息表示訓練已經結束\n        \n    def on_predict_epoch_end(self): #在模型的預測過程中（例如，使用測試數據進行預測）的每個預測訓練周期結束時被呼叫\n        predictions = self.trainer.predict_loop.predictions #取得預測結果、計算生成的圖像的總數、印出生成的圖像數量\n        num_batches = len(predictions)\n        batch_size = predictions[0].shape[0]\n        last_batch_diff = batch_size - predictions[-1].shape[0]\n        print(f\"Number of images generated: {num_batches*batch_size-last_batch_diff}.\")","metadata":{"execution":{"iopub.status.busy":"2023-10-29T01:54:37.979821Z","iopub.execute_input":"2023-10-29T01:54:37.980162Z","iopub.status.idle":"2023-10-29T01:54:38.018988Z","shell.execute_reply.started":"2023-10-29T01:54:37.980131Z","shell.execute_reply":"2023-10-29T01:54:38.018185Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model Training ","metadata":{}},{"cell_type":"code","source":"TRAIN_CONFIG = {  #各種訓練參數的設定\n    \"accelerator\": \"gpu\" if torch.cuda.is_available() else \"cpu\",\n    \n    # train on 16-bit precision\n    \"precision\": \"16-mixed\" if torch.cuda.is_available() else 32,\n    \n    # train on single GPU\n    \"devices\": 1,\n    \n    # save checkpoint only for last epoch by default\n    \"enable_checkpointing\": True,\n    \n    # disable logging for simplicity\n    \"logger\": False,\n    \n    # the number of epochs for training (we limit the number of train/predict batches during debugging)\n    \"max_epochs\": MODEL_CONFIG[\"num_epochs\"],\n    \"limit_train_batches\": 1.0 if not DEBUG else 2,\n    \"limit_predict_batches\": 1.0 if not DEBUG else 5,\n    \n    # the maximum amount of time for training, in case we exceed run-time of 5 hours\n    \"max_time\": {\"hours\": 4, \"minutes\": 55},\n    \n    # use a small subset of photos for validation/testing (we limit here for flexibility)\n    \"limit_val_batches\": 1,\n    \"limit_test_batches\": 5,\n    \n    # disable sanity check before starting the training routine\n    \"num_sanity_val_steps\": 0,\n    \n    # the frequency to visualize the progress of adding Monet style\n    \"check_val_every_n_epoch\": 6 if not DEBUG else 1,\n}","metadata":{"execution":{"iopub.status.busy":"2023-10-29T01:54:38.020096Z","iopub.execute_input":"2023-10-29T01:54:38.020359Z","iopub.status.idle":"2023-10-29T01:54:38.032702Z","shell.execute_reply.started":"2023-10-29T01:54:38.020337Z","shell.execute_reply":"2023-10-29T01:54:38.031985Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"dm = CustomDataModule(**DM_CONFIG)\nmodel = CycleGAN(**MODEL_CONFIG)\ntrainer = L.Trainer(**TRAIN_CONFIG)\ntrainer.fit(model, datamodule=dm)","metadata":{"execution":{"iopub.status.busy":"2023-10-29T01:54:38.033635Z","iopub.execute_input":"2023-10-29T01:54:38.033915Z","iopub.status.idle":"2023-10-29T06:15:25.855553Z","shell.execute_reply.started":"2023-10-29T01:54:38.033892Z","shell.execute_reply":"2023-10-29T06:15:25.854280Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Testing the model on other sample photos**","metadata":{}},{"cell_type":"code","source":"_ = trainer.test(model, datamodule=dm)","metadata":{"execution":{"iopub.status.busy":"2023-10-29T06:15:25.857319Z","iopub.execute_input":"2023-10-29T06:15:25.857649Z","iopub.status.idle":"2023-10-29T06:15:30.352041Z","shell.execute_reply.started":"2023-10-29T06:15:25.857618Z","shell.execute_reply":"2023-10-29T06:15:30.350847Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Submission","metadata":{}},{"cell_type":"markdown","source":"Computing the predictions can be done by running the predict method to generate the Monet-style images given the input photos.","metadata":{}},{"cell_type":"code","source":"predictions = trainer.predict(model, datamodule=dm)","metadata":{"execution":{"iopub.status.busy":"2023-10-29T06:15:30.353717Z","iopub.execute_input":"2023-10-29T06:15:30.354782Z","iopub.status.idle":"2023-10-29T06:16:25.961433Z","shell.execute_reply.started":"2023-10-29T06:15:30.354740Z","shell.execute_reply":"2023-10-29T06:16:25.960346Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Saving the generated images**","metadata":{}},{"cell_type":"code","source":"os.makedirs(\"../images\", exist_ok=True)\nidx = 0\nfor tensor in predictions:\n    for monet in tensor:\n        save_image(\n            monet.float().squeeze() * 0.5 + 0.5, \n            fp=f\"../images/{idx}.jpg\",\n        )\n        idx += 1\n\nshutil.make_archive(\"/kaggle/working/images\", \"zip\", \"/kaggle/images\")","metadata":{"execution":{"iopub.status.busy":"2023-10-29T06:16:25.963388Z","iopub.execute_input":"2023-10-29T06:16:25.963841Z","iopub.status.idle":"2023-10-29T06:16:48.557663Z","shell.execute_reply.started":"2023-10-29T06:16:25.963794Z","shell.execute_reply":"2023-10-29T06:16:48.556603Z"},"trusted":true},"execution_count":null,"outputs":[]}]}