{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"## Purpose\nOriginal repo: [DeepLearning-FromPaper](https://github.com/NinaM31/DeepLearning-FromPaper).    \n\nThis is my implementation of [Unpaired Image-to-Image Translation using Cycle-Consistent Adversarial Networks](https://arxiv.org/pdf/1703.10593.pdf) paper, on the monet problem.","metadata":{}},{"cell_type":"markdown","source":"# Imports","metadata":{}},{"cell_type":"code","source":"import os\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.optim as optim\nfrom torch.utils.data import Dataset, DataLoader\nfrom torchvision import transforms\nfrom PIL import Image\n\nimport numpy as np\nimport pickle as pkl\nimport matplotlib.pyplot as plt\nfrom mpl_toolkits.axes_grid1 import ImageGrid","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2022-05-17T20:50:21.703979Z","iopub.execute_input":"2022-05-17T20:50:21.704601Z","iopub.status.idle":"2022-05-17T20:50:22.323854Z","shell.execute_reply.started":"2022-05-17T20:50:21.704489Z","shell.execute_reply":"2022-05-17T20:50:22.322955Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Custom Dataset","metadata":{}},{"cell_type":"code","source":"class Dataset(torch.utils.data.Dataset):\n\n    def __init__(self, img_dir):\n        img_dir = BASE_DATASET_PATH + \"/\" + img_dir + \"/\"\n        \n        path_list = os.listdir(img_dir)\n        abspath = os.path.abspath(img_dir) \n        \n        self.img_dir = img_dir\n        self.img_list = [os.path.join(abspath, path) for path in path_list]\n\n        self.transform = transforms.Compose([\n            transforms.Resize(IMG_SIZE),\n            transforms.ToTensor(),\n            transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]), # normalize image between -1 and 1\n        ])\n\n\n    def __len__(self):\n        return len(self.img_list)\n\n\n    def __getitem__(self, idx):\n        path = self.img_list[idx]\n        img = Image.open(path).convert('RGB')\n\n        img_tensor = self.transform(img)\n        return img_tensor","metadata":{"execution":{"iopub.status.busy":"2022-05-17T20:50:22.325155Z","iopub.execute_input":"2022-05-17T20:50:22.325385Z","iopub.status.idle":"2022-05-17T20:50:22.333245Z","shell.execute_reply.started":"2022-05-17T20:50:22.325353Z","shell.execute_reply":"2022-05-17T20:50:22.332590Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Discriminator Class","metadata":{}},{"cell_type":"code","source":"class Discriminator(nn.Module):\n\n    def __init__(self,conv_dim=32):\n        super(Discriminator, self).__init__()\n\n        self.main = nn.Sequential(\n            nn.Conv2d(3, conv_dim, 4, stride=2, padding=1),\n            nn.LeakyReLU(0.2, inplace=True),\n\n            nn.Conv2d(conv_dim, conv_dim*2, 4, stride=2, padding=1),\n            nn.InstanceNorm2d(conv_dim*2),\n            nn.LeakyReLU(0.2, inplace=True),\n\n            nn.Conv2d(conv_dim*2, conv_dim*4, 4, stride=2, padding=1),\n            nn.InstanceNorm2d(conv_dim*4),\n            nn.LeakyReLU(0.2, inplace=True),\n\n            nn.Conv2d(conv_dim*4, conv_dim*8, 4, padding=1),\n            nn.InstanceNorm2d(conv_dim*8),\n            nn.LeakyReLU(0.2, inplace=True),\n\n            nn.Conv2d(conv_dim*8, 1, 4, padding=1),\n        )\n\n    def forward(self, x):\n        x = self.main(x)\n        x = F.avg_pool2d(x, x.size()[2:])\n        x = torch.flatten(x, 1)\n        return x","metadata":{"execution":{"iopub.status.busy":"2022-05-17T20:50:22.674218Z","iopub.execute_input":"2022-05-17T20:50:22.674862Z","iopub.status.idle":"2022-05-17T20:50:22.685936Z","shell.execute_reply.started":"2022-05-17T20:50:22.674826Z","shell.execute_reply":"2022-05-17T20:50:22.684795Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# ResidualBlock Class","metadata":{}},{"cell_type":"code","source":"class ResidualBlock(nn.Module):\n    def __init__(self, in_channels):\n        super(ResidualBlock, self).__init__()\n\n        self.main = nn.Sequential(\n            nn.ReflectionPad2d(1),\n            nn.Conv2d(in_channels, in_channels, 3),\n            nn.InstanceNorm2d(in_channels),\n            nn.ReLU(inplace=True),\n            nn.ReflectionPad2d(1),\n            nn.Conv2d(in_channels, in_channels, 3),\n            nn.InstanceNorm2d(in_channels)\n        )\n\n    def forward(self, x):\n        return x + self.main(x)","metadata":{"execution":{"iopub.status.busy":"2022-05-17T20:50:24.839649Z","iopub.execute_input":"2022-05-17T20:50:24.840359Z","iopub.status.idle":"2022-05-17T20:50:24.848011Z","shell.execute_reply.started":"2022-05-17T20:50:24.840321Z","shell.execute_reply":"2022-05-17T20:50:24.847238Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Generator Class","metadata":{}},{"cell_type":"code","source":"class Generator(nn.Module):\n    def __init__(self, conv_dim=64, n_res_block=9):\n        super(Generator, self).__init__()\n        self.main = nn.Sequential(\n            nn.ReflectionPad2d(3),\n            nn.Conv2d(3, conv_dim, 7),\n            nn.InstanceNorm2d(64),\n            nn.ReLU(inplace=True),\n\n            nn.Conv2d(conv_dim, conv_dim*2, 3, stride=2, padding=1),\n            nn.InstanceNorm2d(conv_dim*2),\n            nn.ReLU(inplace=True),\n            nn.Conv2d(conv_dim*2, conv_dim*4, 3, stride=2, padding=1),\n            nn.InstanceNorm2d(conv_dim*4),\n            nn.ReLU(inplace=True),\n\n            ResidualBlock(conv_dim*4),\n            ResidualBlock(conv_dim*4),\n            ResidualBlock(conv_dim*4),\n            ResidualBlock(conv_dim*4),\n            ResidualBlock(conv_dim*4),\n            ResidualBlock(conv_dim*4),\n            ResidualBlock(conv_dim*4),\n            ResidualBlock(conv_dim*4),\n            ResidualBlock(conv_dim*4),\n\n            nn.ConvTranspose2d(conv_dim*4, conv_dim*2, 3, stride=2, padding=1, output_padding=1),\n            nn.InstanceNorm2d(conv_dim*2),\n            nn.ReLU(inplace=True),\n            nn.ConvTranspose2d(conv_dim*2, conv_dim, 3, stride=2, padding=1, output_padding=1),\n            nn.InstanceNorm2d(conv_dim),\n            nn.ReLU(inplace=True),\n\n            nn.ReflectionPad2d(3),\n            nn.Conv2d(conv_dim, 3, 7),\n            nn.Tanh()\n        )\n\n    def forward(self, x):\n        return self.main(x)","metadata":{"execution":{"iopub.status.busy":"2022-05-17T20:50:26.014620Z","iopub.execute_input":"2022-05-17T20:50:26.015319Z","iopub.status.idle":"2022-05-17T20:50:26.026547Z","shell.execute_reply.started":"2022-05-17T20:50:26.015282Z","shell.execute_reply":"2022-05-17T20:50:26.025463Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# CycleGAN Class","metadata":{}},{"cell_type":"code","source":"class CycleGAN:\n\n    def __init__(self, g_conv_dim=64, d_conv_dim=64, n_res_block=6):\n        self.device = torch.device('cuda') if torch.cuda.is_available() else torch.device(\"cpu\")\n\n        self.G_XtoY = Generator(conv_dim=g_conv_dim, n_res_block=n_res_block).to(self.device)\n        self.G_YtoX = Generator(conv_dim=g_conv_dim, n_res_block=n_res_block).to(self.device)\n\n        self.D_X = Discriminator(conv_dim=d_conv_dim).to(self.device)\n        self.D_Y = Discriminator(conv_dim=d_conv_dim).to(self.device)\n\n        print(f\"Models running of {self.device}\")\n\n    def load_model(self, filename):\n        save_filename = os.path.splitext(os.path.basename(filename))[0] + '.pt'\n        return torch.load(save_filename)\n\n    def real_mse_loss(self, D_out):\n        return torch.mean((D_out-1)**2)\n\n\n    def fake_mse_loss(self, D_out):\n        return torch.mean(D_out**2)\n\n\n    def cycle_consistency_loss(self, real_img, reconstructed_img, lambda_weight):\n        reconstr_loss = torch.mean(torch.abs(real_img - reconstructed_img))\n        return lambda_weight*reconstr_loss    \n\n    \n    def train_generator(self, optimizers, images_x, images_y):\n        # Generator YtoX\n        optimizers[\"g_optim\"].zero_grad()\n\n        fake_images_x = self.G_YtoX(images_y)\n\n        d_real_x = self.D_X(fake_images_x)\n        g_YtoX_loss = self.real_mse_loss(d_real_x)\n\n        recon_y = self.G_XtoY(fake_images_x)\n        recon_y_loss = self.cycle_consistency_loss(images_y, recon_y, lambda_weight=10)\n\n\n        # Generator XtoY\n        fake_images_y = self.G_XtoY(images_x)\n\n        d_real_y = self.D_Y(fake_images_y)\n        g_XtoY_loss = self.real_mse_loss(d_real_y)\n\n        recon_x = self.G_YtoX(fake_images_y)\n        recon_x_loss = self.cycle_consistency_loss(images_x, recon_x, lambda_weight=10)\n\n        g_total_loss = g_YtoX_loss + g_XtoY_loss + recon_y_loss + recon_x_loss\n        g_total_loss.backward()\n        optimizers[\"g_optim\"].step()\n\n        return g_total_loss.item()\n\n    \n    def train_discriminator(self, optimizers, images_x, images_y):\n        # Discriminator x\n        optimizers[\"d_x_optim\"].zero_grad()\n\n        d_real_x = self.D_X(images_x)\n        d_real_loss_x = self.real_mse_loss(d_real_x)\n        \n        fake_images_x = self.G_YtoX(images_y)\n\n        d_fake_x = self.D_X(fake_images_x)\n        d_fake_loss_x = self.fake_mse_loss(d_fake_x)\n        \n        d_x_loss = d_real_loss_x + d_fake_loss_x\n        d_x_loss.backward()\n        optimizers[\"d_x_optim\"].step()\n\n\n        # Discriminator y\n        optimizers[\"d_y_optim\"].zero_grad()\n            \n        d_real_y = self.D_Y(images_y)\n        d_real_loss_x = self.real_mse_loss(d_real_y)\n    \n        fake_images_y = self.G_XtoY(images_x)\n\n        d_fake_y = self.D_Y(fake_images_y)\n        d_fake_loss_y = self.fake_mse_loss(d_fake_y)\n\n        d_y_loss = d_real_loss_x + d_fake_loss_y\n        d_y_loss.backward()\n        optimizers[\"d_y_optim\"].step()\n\n        return d_x_loss.item(), d_y_loss.item()\n\n\n    def train(self, optimizers, data_loader_x, data_loader_y, print_every=1, sample_every=100):\n        losses = []\n        g_total_loss_min = np.Inf\n    \n        fixed_x = next(iter(data_loader_x))[1].to(self.device)\n        fixed_y = next(iter(data_loader_y))[1].to(self.device)\n\n        print(f'Running on {self.device}')\n        for epoch in range(EPOCHS):\n            for (images_x, images_y) in zip(data_loader_x, data_loader_y):\n                images_x, images_y = images_x.to(self.device), images_y.to(self.device)\n                \n                g_total_loss = self.train_generator(optimizers, images_x, images_y)\n                d_x_loss, d_y_loss = self.train_discriminator(optimizers, images_x, images_y)\n                \n            \n            if epoch % print_every == 0:\n                losses.append((d_x_loss, d_y_loss, g_total_loss))\n                print('Epoch [{:5d}/{:5d}] | d_X_loss: {:6.4f} | d_Y_loss: {:6.4f} | g_total_loss: {:6.4f}'\n                .format(\n                    epoch, \n                    EPOCHS, \n                    d_x_loss, \n                    d_y_loss, \n                    g_total_loss\n                ))\n                \n            if g_total_loss < g_total_loss_min:\n                g_total_loss_min = g_total_loss\n                \n                torch.save(self.G_XtoY.state_dict(), \"G_X2Y\")\n                torch.save(self.G_YtoX.state_dict(), \"G_Y2X\")\n                \n                torch.save(self.D_X.state_dict(), \"D_X\")\n                torch.save(self.D_Y.state_dict(), \"D_Y\")\n                \n                print(\"Models Saved\")\n                \n                \n\n        return losses","metadata":{"execution":{"iopub.status.busy":"2022-05-17T20:50:28.227185Z","iopub.execute_input":"2022-05-17T20:50:28.227473Z","iopub.status.idle":"2022-05-17T20:50:28.252188Z","shell.execute_reply.started":"2022-05-17T20:50:28.227441Z","shell.execute_reply":"2022-05-17T20:50:28.251452Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Config","metadata":{}},{"cell_type":"code","source":"BASE_DATASET_PATH = \"../input/gan-getting-started\"\nX_DATASET = \"photo_jpg\"\nY_DATASET = \"monet_jpg\"\n\nBATCH_SIZE = 32\nN_WORKERS = 0\n\nIMG_SIZE = 128\nLR = 0.0002\nBETA1 = 0.5\nBETA2 = 0.999\n\nEPOCHS = 50","metadata":{"execution":{"iopub.status.busy":"2022-05-17T20:50:37.535006Z","iopub.execute_input":"2022-05-17T20:50:37.535265Z","iopub.status.idle":"2022-05-17T20:50:37.539612Z","shell.execute_reply.started":"2022-05-17T20:50:37.535236Z","shell.execute_reply":"2022-05-17T20:50:37.538935Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Train","metadata":{}},{"cell_type":"code","source":"# Dataset\nx_dataset = Dataset(X_DATASET)\ny_dataset = Dataset(Y_DATASET)\n\ndata_loader_x = DataLoader(x_dataset, BATCH_SIZE, shuffle=True, num_workers=N_WORKERS)\ndata_loader_y = DataLoader(y_dataset, BATCH_SIZE, shuffle=True, num_workers=N_WORKERS)\n\n# Model\ncycleGan = CycleGAN()\n\n# Oprimizer\ng_params = list(cycleGan.G_XtoY.parameters()) + list(cycleGan.G_YtoX.parameters())\n\noptimizers = {\n    \"g_optim\": optim.Adam(g_params, LR, [BETA1, BETA2]),\n    \"d_x_optim\": optim.Adam(cycleGan.D_X.parameters(), LR, [BETA1, BETA2]),\n    \"d_y_optim\": optim.Adam(cycleGan.D_Y.parameters(), LR, [BETA1, BETA2])\n}\n\n# Train\nlosses = cycleGan.train(optimizers, data_loader_x, data_loader_y)","metadata":{"execution":{"iopub.status.busy":"2022-05-17T20:51:33.337135Z","iopub.execute_input":"2022-05-17T20:51:33.337425Z","iopub.status.idle":"2022-05-17T21:12:41.141293Z","shell.execute_reply.started":"2022-05-17T20:51:33.337380Z","shell.execute_reply":"2022-05-17T21:12:41.140527Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Results","metadata":{}},{"cell_type":"code","source":"# Plot\nfig, ax = plt.subplots(figsize=(12,8))\nlosses = np.array(losses)\nplt.plot(losses.T[0], label='Discriminator, X', alpha=0.5)\nplt.plot(losses.T[1], label='Discriminator, Y', alpha=0.5)\nplt.plot(losses.T[2], label='Generators', alpha=0.5)\nplt.title(\"Training Losses\")\nplt.legend()\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2022-05-17T21:12:41.142841Z","iopub.execute_input":"2022-05-17T21:12:41.143431Z","iopub.status.idle":"2022-05-17T21:12:41.361871Z","shell.execute_reply.started":"2022-05-17T21:12:41.143336Z","shell.execute_reply":"2022-05-17T21:12:41.361225Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"samples = []\n\nfor i in range(12):\n    fixed_x = next(iter(data_loader_x))[i].to(cycleGan.device)\n    fake_y = cycleGan.G_XtoY(torch.unsqueeze(fixed_x, dim=0))\n    samples.extend([fixed_x, torch.squeeze(fake_y, 0)])","metadata":{"execution":{"iopub.status.busy":"2022-05-17T21:12:41.363124Z","iopub.execute_input":"2022-05-17T21:12:41.363369Z","iopub.status.idle":"2022-05-17T21:12:42.524576Z","shell.execute_reply.started":"2022-05-17T21:12:41.363319Z","shell.execute_reply":"2022-05-17T21:12:42.523720Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"fig = plt.figure(figsize=(18, 14))\ngrid = ImageGrid(fig, 111, nrows_ncols=(2, 4), axes_pad=0.5)\n\n\nfor i, (ax, im) in enumerate(zip(grid, samples)):\n    _, w, h = im.size()\n    im = im.detach().cpu().numpy()\n    im = np.transpose(im, (1, 2, 0))\n    \n    im = ((im +1)*255 / (2)).astype(np.uint8)\n    ax.imshow(im.reshape((w,h,3)))\n\n    ax.xaxis.set_visible(False)\n    ax.yaxis.set_visible(False)\n\n    if i%2 == 0: title = \"Original\"\n    else: title = \"fake\"\n\n    ax.set_title(title)\n\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2022-05-17T21:15:17.819985Z","iopub.execute_input":"2022-05-17T21:15:17.820274Z","iopub.status.idle":"2022-05-17T21:15:18.924097Z","shell.execute_reply.started":"2022-05-17T21:15:17.820243Z","shell.execute_reply":"2022-05-17T21:15:18.923267Z"},"trusted":true},"execution_count":null,"outputs":[]}]}