{"cells":[{"metadata":{"trusted":true},"cell_type":"code","source":"DEBUG = False","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import os\nimport pandas as pd\nimport numpy as np\nfrom PIL import Image\nimport matplotlib.pyplot as plt\nimport itertools\nfrom tqdm.notebook import tqdm\nimport wandb\n\nimport torch\nfrom torch.utils.data import Dataset, DataLoader\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torchvision\nfrom torchvision import transforms","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!wandb login 45e03af61cb9a1b88a119d38b52d0396b3d33437","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"NAME = 'debug' if DEBUG else 'cycle-gan-exp01'\n\nconfig = dict( imgsize = 256,\n             batch_size = 1 if DEBUG else 2,\n             input_nc = 3,\n             output_nc = 3,\n             lr = 2e-4,\n             epoch = 0,\n             n_epochs = 30,\n             decay_epoch = 15)\n\nwandb.init(project = 'cyclegan', name = NAME, config = config)","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"IMG_PATH = '../input/gan-getting-started/photo_jpg'\nMONET_PATH = '../input/gan-getting-started/monet_jpg'\n\nimg_list = os.listdir(IMG_PATH)\nmonet_list = os.listdir(MONET_PATH)\n\nprint(len(img_list))\nprint(len(monet_list))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plt.figure(figsize = (15,10))\n\nfor i in range(10):\n    plt.subplot(2,5,i+1)\n    img = Image.open(os.path.join(IMG_PATH,img_list[i])).convert('RGB')\n    plt.imshow(img)\n    plt.axis('off')\n    \nplt.tight_layout()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plt.figure(figsize = (15,10))\n\nfor i in range(10):\n    plt.subplot(2,5,i+1)\n    img = Image.open(os.path.join(MONET_PATH,monet_list[i])).convert('RGB')\n    plt.imshow(img)\n    plt.axis('off')\n    \nplt.tight_layout()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class IMGDataset(Dataset):\n    def __init__(self, config, tfm):\n        self.size = config['imgsize']\n        self.img_list = os.listdir(IMG_PATH)\n        self.monet_list = os.listdir(MONET_PATH)\n        self.tfm = tfm\n        \n    def __len__(self):\n        return min(len(self.img_list), len(self.monet_list))\n    \n    def __getitem__(self, idx):\n        \n#         img_index = idx % len(self.img_list)\n#         monet_index = idx % len(self.monet_list)\n\n        img_index = idx\n        monet_index = idx\n        \n        img = Image.open(os.path.join(IMG_PATH, self.img_list[img_index])).convert('RGB')\n        monet_img = Image.open(os.path.join(MONET_PATH, self.monet_list[monet_index])).convert('RGB')\n        \n        img = self.tfm(img)\n        monet_img = self.tfm(monet_img)\n        \n        return img, monet_img","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"tfm =  transforms.Compose([transforms.ToTensor(),\n        transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5)) ])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"dataset = IMGDataset(config, tfm)\ndataloader = DataLoader(dataset, batch_size = config['batch_size'], shuffle = True, num_workers = True, pin_memory=True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class ResidualBlock(nn.Module):\n    def __init__(self, in_features):\n        super(ResidualBlock, self).__init__()\n\n        conv_block = [  nn.ReflectionPad2d(1),\n                        nn.Conv2d(in_features, in_features, 3),\n                        nn.InstanceNorm2d(in_features),\n                        nn.ReLU(inplace=True),\n                        nn.ReflectionPad2d(1),\n                        nn.Conv2d(in_features, in_features, 3),\n                        nn.InstanceNorm2d(in_features)  ]\n\n        self.conv_block = nn.Sequential(*conv_block)\n\n    def forward(self, x):\n        return x + self.conv_block(x)\n\nclass Generator(nn.Module):\n    def __init__(self, input_nc, output_nc, n_residual_blocks=9):\n        super(Generator, self).__init__()\n\n        # Initial convolution block       \n        model = [   nn.ReflectionPad2d(3),\n                    nn.Conv2d(input_nc, 64, 7),\n                    nn.InstanceNorm2d(64),\n                    nn.ReLU(inplace=True) ]\n\n        # Downsampling\n        in_features = 64\n        out_features = in_features*2\n        for _ in range(2):\n            model += [  nn.Conv2d(in_features, out_features, 3, stride=2, padding=1),\n                        nn.InstanceNorm2d(out_features),\n                        nn.ReLU(inplace=True) ]\n            in_features = out_features\n            out_features = in_features*2\n\n        # Residual blocks\n        for _ in range(n_residual_blocks):\n            model += [ResidualBlock(in_features)]\n\n        # Upsampling\n        out_features = in_features//2\n        for _ in range(2):\n            model += [  nn.ConvTranspose2d(in_features, out_features, 3, stride=2, padding=1, output_padding=1),\n                        nn.InstanceNorm2d(out_features),\n                        nn.ReLU(inplace=True) ]\n            in_features = out_features\n            out_features = in_features//2\n\n        # Output layer\n        model += [  nn.ReflectionPad2d(3),\n                    nn.Conv2d(64, output_nc, 7),\n                    nn.Tanh() ]\n\n        self.model = nn.Sequential(*model)\n\n    def forward(self, x):\n        return self.model(x)\n\nclass Discriminator(nn.Module):\n    def __init__(self, input_nc):\n        super(Discriminator, self).__init__()\n\n        # A bunch of convolutions one after another\n        model = [   nn.Conv2d(input_nc, 64, 4, stride=2, padding=1),\n                    nn.LeakyReLU(0.2, inplace=True) ]\n\n        model += [  nn.Conv2d(64, 128, 4, stride=2, padding=1),\n                    nn.InstanceNorm2d(128), \n                    nn.LeakyReLU(0.2, inplace=True) ]\n\n        model += [  nn.Conv2d(128, 256, 4, stride=2, padding=1),\n                    nn.InstanceNorm2d(256), \n                    nn.LeakyReLU(0.2, inplace=True) ]\n\n        model += [  nn.Conv2d(256, 512, 4, padding=1),\n                    nn.InstanceNorm2d(512), \n                    nn.LeakyReLU(0.2, inplace=True) ]\n\n        # FCN classification layer\n        model += [nn.Conv2d(512, 1, 4, padding=1)]\n\n        self.model = nn.Sequential(*model)\n\n    def forward(self, x):\n        x =  self.model(x)\n        # Average pooling and flatten\n        return F.avg_pool2d(x, x.size()[2:]).view(x.size()[0], -1)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def weights_init_normal(m):\n    classname = m.__class__.__name__\n    if classname.find('Conv') != -1:\n        torch.nn.init.normal(m.weight.data, 0.0, 0.02)\n    elif classname.find('BatchNorm2d') != -1:\n        torch.nn.init.normal(m.weight.data, 1.0, 0.02)\n        torch.nn.init.constant(m.bias.data, 0.0)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class LambdaLR():\n    def __init__(self, n_epochs, offset, decay_start_epoch):\n        assert ((n_epochs - decay_start_epoch) > 0), \"Decay must start before the training session ends!\"\n        self.n_epochs = n_epochs\n        self.offset = offset\n        self.decay_start_epoch = decay_start_epoch\n\n    def step(self, epoch):\n        return 1.0 - max(0, epoch + self.offset - self.decay_start_epoch)/(self.n_epochs - self.decay_start_epoch)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"netG_A2B = Generator(config['input_nc'], config['output_nc'])\nnetG_B2A = Generator(config['input_nc'], config['output_nc'])\nnetD_A = Discriminator(config['input_nc'])\nnetD_B = Discriminator(config['output_nc'])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"netG_A2B.apply(weights_init_normal)\nnetG_B2A.apply(weights_init_normal)\nnetD_A.apply(weights_init_normal)\nnetD_B.apply(weights_init_normal)\n\nprint('Weights Initialized')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"device = 'cuda' if torch.cuda.is_available() else 'cpu'\n\nnetG_A2B.to(device)\nnetG_B2A.to(device)\nnetD_A.to(device) \nnetD_B.to(device) \n\nprint(f'Transferred to {device}')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Lossess\ncriterion_GAN = torch.nn.MSELoss()\ncriterion_cycle = torch.nn.L1Loss()\ncriterion_identity = torch.nn.L1Loss()\n\n# Optimizers & LR schedulers\noptimizer_G = torch.optim.Adam(itertools.chain(netG_A2B.parameters(), netG_B2A.parameters()),\n                                lr=config['lr'], betas=(0.5, 0.999))\noptimizer_D_A = torch.optim.Adam(netD_A.parameters(), lr=config['lr'], betas=(0.5, 0.999))\noptimizer_D_B = torch.optim.Adam(netD_B.parameters(), lr=config['lr'], betas=(0.5, 0.999))\n\nlr_scheduler_G = torch.optim.lr_scheduler.LambdaLR(optimizer_G, lr_lambda=LambdaLR(config['n_epochs'], config['epoch'], config['decay_epoch']).step)\nlr_scheduler_D_A = torch.optim.lr_scheduler.LambdaLR(optimizer_D_A, lr_lambda=LambdaLR(config['n_epochs'], config['epoch'], config['decay_epoch']).step)\nlr_scheduler_D_B = torch.optim.lr_scheduler.LambdaLR(optimizer_D_B, lr_lambda=LambdaLR(config['n_epochs'], config['epoch'], config['decay_epoch']).step)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"target_real = torch.ones(config['batch_size'], dtype=torch.float).unsqueeze(1).to(device)\ntarget_fake = torch.ones(config['batch_size'], dtype=torch.float).unsqueeze(1).to(device)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"wandb_step = 0\nlog_image_step = 50\nloader_len = len(dataloader)\n\n###### Training ######\nfor epoch in range(config['epoch'], config['n_epochs']):\n    for i, (photo, monet_img) in enumerate(tqdm(dataloader, total = loader_len)):\n        \n        # Set model input\n        real_A = photo.to(device)\n        real_B = monet_img.to(device)\n\n        ###### Generators A2B and B2A ######\n        optimizer_G.zero_grad()\n\n        # Identity loss\n        # G_A2B(B) should equal B if real B is fed\n        same_B = netG_A2B(real_B)\n        loss_identity_B = criterion_identity(same_B, real_B)*5.0\n        # G_B2A(A) should equal A if real A is fed\n        same_A = netG_B2A(real_A)\n        loss_identity_A = criterion_identity(same_A, real_A)*5.0\n\n        # GAN loss\n        fake_B = netG_A2B(real_A)\n        pred_fake = netD_B(fake_B)\n        loss_GAN_A2B = criterion_GAN(pred_fake, target_real)\n\n        fake_A = netG_B2A(real_B)\n        pred_fake = netD_A(fake_A)\n        loss_GAN_B2A = criterion_GAN(pred_fake, target_real)\n\n        # Cycle loss\n        recovered_A = netG_B2A(fake_B)\n        loss_cycle_ABA = criterion_cycle(recovered_A, real_A)*10.0\n\n        recovered_B = netG_A2B(fake_A)\n        loss_cycle_BAB = criterion_cycle(recovered_B, real_B)*10.0\n\n        # Total loss\n        loss_G = loss_identity_A + loss_identity_B + loss_GAN_A2B + loss_GAN_B2A + loss_cycle_ABA + loss_cycle_BAB\n        loss_G.backward(retain_graph = True)\n        \n        \n        ###################################\n\n        ###### Discriminator A ######\n        optimizer_D_A.zero_grad()\n\n        # Real loss\n        pred_real = netD_A(real_A)\n        loss_D_real = criterion_GAN(pred_real, target_real)\n\n        # Fake loss\n        pred_fake = netD_A(fake_A)\n        loss_D_fake = criterion_GAN(pred_fake, target_fake)\n\n        # Total loss\n        loss_D_A = (loss_D_real + loss_D_fake)*0.5\n        loss_D_A.backward()\n\n        \n        ###################################\n\n        ###### Discriminator B ######\n        optimizer_D_B.zero_grad()\n\n        # Real loss\n        pred_real = netD_B(real_B)\n        loss_D_real = criterion_GAN(pred_real, target_real)\n        \n        # Fake loss\n        pred_fake = netD_B(fake_B)\n        loss_D_fake = criterion_GAN(pred_fake, target_fake)\n\n        # Total loss\n        loss_D_B = (loss_D_real + loss_D_fake)*0.5\n        loss_D_B.backward()\n\n        \n        ###################################\n        \n        # Update the weights\n        optimizer_G.step()\n        optimizer_D_A.step()\n        optimizer_D_B.step()\n        \n        \n\n#         # Progress report (http://localhost:8097)\n\n        wandb_step = loader_len*epoch + i\n        wandb.log({'loss_G': loss_G, 'loss_G_identity': (loss_identity_A + loss_identity_B), 'loss_G_GAN': (loss_GAN_A2B + loss_GAN_B2A),\n                    'loss_G_cycle': (loss_cycle_ABA + loss_cycle_BAB), 'loss_D': (loss_D_A + loss_D_B)}, step = wandb_step)\n        \n        if wandb_step % log_image_step == 0:\n    \n            wandb.log({'exp01': [wandb.Image(real_A.cpu().detach().numpy()[0].transpose(1,2,0), caption='real_A'), \n                                wandb.Image(real_B.cpu().detach().numpy()[0].transpose(1,2,0), caption='real_B'),\n                                wandb.Image(fake_A.cpu().detach().numpy()[0].transpose(1,2,0), caption = 'fake_A'),\n                                wandb.Image(fake_B.cpu().detach().numpy()[0].transpose(1,2,0), caption = 'fake_B')]}, step = wandb_step)\n\n\n\n    # Update learning rates\n    lr_scheduler_G.step()\n    lr_scheduler_D_A.step()\n    lr_scheduler_D_B.step()","execution_count":null,"outputs":[]}],"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":4,"nbformat_minor":4}