{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":21755,"databundleVersionId":1475600,"sourceType":"competition"}],"dockerImageVersionId":30528,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import glob\nimport random\nimport os\nimport time\nimport datetime\nimport sys\nimport math\nimport itertools\nimport shutil\n\n\nimport numpy as np\nfrom PIL import Image\nimport matplotlib.pyplot as plt\n\n\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.cuda.amp import autocast\n\nfrom torch.utils.data import Dataset, DataLoader\nfrom torchvision.transforms import ToTensor, Normalize, Compose, Resize, RandomCrop, RandomHorizontalFlip\nfrom torchvision.io import read_image\nfrom torchvision.utils import save_image, make_grid\nfrom torch.autograd import Variable","metadata":{"execution":{"iopub.status.busy":"2023-08-10T09:18:38.406262Z","iopub.execute_input":"2023-08-10T09:18:38.406581Z","iopub.status.idle":"2023-08-10T09:18:42.446061Z","shell.execute_reply.started":"2023-08-10T09:18:38.406554Z","shell.execute_reply":"2023-08-10T09:18:42.444965Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"root = \"/kaggle/input/gan-getting-started/\"\n\nBATCH_SIZE = 12\n\ntransforms = [\n    Resize(int(256 * 1.12), Image.BICUBIC),\n    RandomCrop((256, 256)),\n    RandomHorizontalFlip(),\n    ToTensor(),\n    Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))\n]\n\nclass ImageDataset(Dataset):\n    def __init__(self, root, transforms, mode='train', unaligned=True):\n        self.mode = mode\n        self.unaligned = unaligned\n        self.transform = Compose(transforms)\n        self.files_A = sorted(glob.glob(os.path.join(root+'/monet_jpg')+'/*.*'))\n        self.files_B = sorted(glob.glob(os.path.join(root+'/photo_jpg')+'/*.*'))\n\n    def  __getitem__(self, index):\n        image_A = Image.open(self.files_A[index % len(self.files_A)])\n        \n        if self.unaligned:\n            image_B = Image.open(self.files_B[torch.randint(0, len(self.files_B)-1, (1,1))])\n        else:\n            image_B = Image.open(self.files_B[index % len(self.files_B)])\n            \n        item_A = self.transform(image_A)\n        item_B = self.transform(image_B)\n        return {'A':item_A, 'B':item_B}\n    \n    def __len__(self):\n        if self.mode == \"train\":\n            return min(len(self.files_A), len(self.files_B))\n        else:\n            return max(len(self.files_A), len(self.files_B))\n\ndataloader = DataLoader(\n    ImageDataset(root, transforms=transforms, unaligned=True),\n    batch_size=BATCH_SIZE,\n    shuffle=True,\n    num_workers=os.cpu_count(),\n    pin_memory = torch.cuda.is_available()\n)\n\nval_transforms = transforms = [\n    ToTensor(),\n    Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))\n]\n\n\nval_dataloader = DataLoader(\n    ImageDataset(root, transforms=val_transforms, mode='gen', unaligned=False),\n    batch_size=1,\n    shuffle=False,\n    num_workers=os.cpu_count(),\n    pin_memory = torch.cuda.is_available()\n)","metadata":{"execution":{"iopub.status.busy":"2023-08-10T09:18:42.45138Z","iopub.execute_input":"2023-08-10T09:18:42.452734Z","iopub.status.idle":"2023-08-10T09:18:42.830493Z","shell.execute_reply.started":"2023-08-10T09:18:42.452697Z","shell.execute_reply":"2023-08-10T09:18:42.828478Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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        if hasattr(m, \"bias\") and m.bias is not None:\n            torch.nn.init.constant_(m.bias.data, 0.0)\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)","metadata":{"execution":{"iopub.status.busy":"2023-08-10T09:18:42.845268Z","iopub.execute_input":"2023-08-10T09:18:42.849983Z","iopub.status.idle":"2023-08-10T09:18:42.860838Z","shell.execute_reply.started":"2023-08-10T09:18:42.849947Z","shell.execute_reply":"2023-08-10T09:18:42.859942Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class ResidualBlock(nn.Module):\n    def __init__(self, in_features):\n        super(ResidualBlock, self).__init__()\n\n        self.block = nn.Sequential(\n            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    #@autocast()\n    def forward(self, x):\n        return x + self.block(x)\n\n\nclass GeneratorResNet(nn.Module):\n    def __init__(self, input_shape, num_residual_blocks):\n        super(GeneratorResNet, self).__init__()\n\n        channels = input_shape[0]\n\n        # Initial convolution block\n        out_features = 64\n        model = [\n            nn.ReflectionPad2d(channels),\n            nn.Conv2d(channels, out_features, 7),\n            nn.InstanceNorm2d(out_features),\n            nn.ReLU(inplace=True),\n        ]\n        in_features = out_features\n\n        # Downsampling\n        for _ in range(2):\n            out_features *= 2\n            model += [\n                nn.Conv2d(in_features, out_features, 3, stride=2, padding=1),\n                nn.InstanceNorm2d(out_features),\n                nn.ReLU(inplace=True),\n            ]\n            in_features = out_features\n\n        # Residual blocks\n        for _ in range(num_residual_blocks):\n            model += [ResidualBlock(out_features)]\n\n        # Upsampling\n        for _ in range(2):\n            out_features //= 2\n            model += [\n                nn.Upsample(scale_factor=2),\n                nn.Conv2d(in_features, out_features, 3, stride=1, padding=1),\n                nn.InstanceNorm2d(out_features),\n                nn.ReLU(inplace=True),\n            ]\n            in_features = out_features\n\n        # Output layer\n        model += [nn.ReflectionPad2d(channels), nn.Conv2d(out_features, channels, 7), nn.Tanh()]\n\n        self.model = nn.Sequential(*model)\n    \n    #@autocast()\n    def forward(self, x):\n        return self.model(x)","metadata":{"execution":{"iopub.status.busy":"2023-08-10T09:18:42.865746Z","iopub.execute_input":"2023-08-10T09:18:42.86781Z","iopub.status.idle":"2023-08-10T09:18:42.892821Z","shell.execute_reply.started":"2023-08-10T09:18:42.867765Z","shell.execute_reply":"2023-08-10T09:18:42.891778Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Discriminator(nn.Module):\n    def __init__(self, input_shape):\n        super(Discriminator, self).__init__()\n\n        channels, height, width = input_shape\n\n        # Calculate output shape of image discriminator (PatchGAN)\n        self.output_shape = (1, height // 2 ** 4, width // 2 ** 4)\n\n        def discriminator_block(in_filters, out_filters, normalize=True):\n            \"\"\"Returns downsampling layers of each discriminator block\"\"\"\n            layers = [nn.Conv2d(in_filters, out_filters, 4, stride=2, padding=1)]\n            if normalize:\n                layers.append(nn.InstanceNorm2d(out_filters))\n            layers.append(nn.LeakyReLU(0.2, inplace=True))\n            return layers\n\n        self.model = nn.Sequential(\n            *discriminator_block(channels, 64, normalize=False),\n            *discriminator_block(64, 128),\n            *discriminator_block(128, 256),\n            *discriminator_block(256, 512),\n            nn.ZeroPad2d((1, 0, 1, 0)),\n            nn.Conv2d(512, 1, 4, padding=1)\n        )\n    \n    #@autocast()\n    def forward(self, img):\n        return self.model(img)","metadata":{"execution":{"iopub.status.busy":"2023-08-10T09:18:42.898947Z","iopub.execute_input":"2023-08-10T09:18:42.902864Z","iopub.status.idle":"2023-08-10T09:18:42.917585Z","shell.execute_reply.started":"2023-08-10T09:18:42.902829Z","shell.execute_reply":"2023-08-10T09:18:42.916455Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"epoch = 0\nn_epochs = 200\ndecay_epoch = int(n_epochs//2)\nbatch_size = 1\nlr = 0.0002\nb1 = 0.5\nb2 = 0.999\nn_cpu = os.cpu_count()\nimg_height = 256\nimg_width = 256\nchannels = 3\nsample_interval = 5\ncheckpoint_interval = 5\nn_residual_blocks = 9\nlambda_cyc = 10.0\nlambda_id = 5.0","metadata":{"execution":{"iopub.status.busy":"2023-08-10T09:18:42.923183Z","iopub.execute_input":"2023-08-10T09:18:42.926835Z","iopub.status.idle":"2023-08-10T09:18:42.935317Z","shell.execute_reply.started":"2023-08-10T09:18:42.926801Z","shell.execute_reply":"2023-08-10T09:18:42.934444Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class ReplayBuffer:\n    def __init__(self, max_size=50):\n        assert max_size > 0, \"Empty buffer or trying to create a black hole. Be careful.\"\n        self.max_size = max_size\n        self.data = []\n\n    def push_and_pop(self, data):\n        to_return = []\n        for element in data.data:\n            element = torch.unsqueeze(element, 0)\n            if len(self.data) < self.max_size:\n                self.data.append(element)\n                to_return.append(element)\n            else:\n                if random.uniform(0, 1) > 0.5:\n                    i = random.randint(0, self.max_size - 1)\n                    to_return.append(self.data[i].clone())\n                    self.data[i] = element\n                else:\n                    to_return.append(element)\n        return Variable(torch.cat(to_return))\n\n\nclass 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)","metadata":{"execution":{"iopub.status.busy":"2023-08-10T09:18:42.941463Z","iopub.execute_input":"2023-08-10T09:18:42.942828Z","iopub.status.idle":"2023-08-10T09:18:42.959507Z","shell.execute_reply.started":"2023-08-10T09:18:42.942797Z","shell.execute_reply":"2023-08-10T09:18:42.958458Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def init_weights(net, init_gain=0.02):\n    def init_func(m):  \n        classname = m.__class__.__name__\n        if classname.find(\"Conv\") != -1:\n            nn.init.normal_(m.weight.data, 0.0, init_gain)\n            if hasattr(m, \"bias\") and m.bias is not None:\n                torch.nn.init.constant_(m.bias.data, 0.0)\n    net.apply(init_func) \n\n\ndef init_net(net, init_gain=0.02, gpu_ids=[]):\n    if len(gpu_ids) > 0:\n        assert(torch.cuda.is_available())\n        net.to(gpu_ids[0])\n        net = torch.nn.DataParallel(net, gpu_ids)  # multi-GPUs\n    init_weights(net, init_gain=init_gain)\n    return net","metadata":{"execution":{"iopub.status.busy":"2023-08-10T09:18:42.96624Z","iopub.execute_input":"2023-08-10T09:18:42.969341Z","iopub.status.idle":"2023-08-10T09:18:42.981319Z","shell.execute_reply.started":"2023-08-10T09:18:42.969308Z","shell.execute_reply":"2023-08-10T09:18:42.980204Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def sample_images(batches_done, print_image=True, batch_size=12):\n    imgs = next(iter(dataloader))\n    G_AB.eval()\n    G_BA.eval()\n    real_A = Variable(imgs[\"A\"]).cuda()\n    fake_B = G_AB(real_A)\n    real_B = Variable(imgs[\"B\"]).cuda()\n    fake_A = G_BA(real_B)\n    real_A = make_grid(real_A, nrow=batch_size//2, normalize=True)\n    real_B = make_grid(real_B, nrow=batch_size//2, normalize=True)\n    fake_A = make_grid(fake_A, nrow=batch_size//2, normalize=True)\n    fake_B = make_grid(fake_B, nrow=batch_size//2, normalize=True)\n    image_grid = torch.cat((real_A, fake_B, real_B, fake_A), 1)\n    save_image(image_grid, \"images_train/%s.png\" % batches_done, normalize=False)\n    if print_image:\n        image_grid = image_grid.permute(1,2,0).cpu().detach().numpy()\n        plt.figure(figsize=(20, 20))\n        plt.imshow(image_grid)\n        plt.show()","metadata":{"execution":{"iopub.status.busy":"2023-08-10T09:18:42.989989Z","iopub.execute_input":"2023-08-10T09:18:42.993344Z","iopub.status.idle":"2023-08-10T09:18:43.006479Z","shell.execute_reply.started":"2023-08-10T09:18:42.993311Z","shell.execute_reply":"2023-08-10T09:18:43.005446Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"os.makedirs(\"images_train\", exist_ok=True)\nos.makedirs(\"saved_models\", exist_ok=True)\n\ncuda = torch.cuda.is_available()\nif torch.cuda.device_count() > 1:\n    gpu_ids = [0,1]\n\n# Losses\ncriterion_GAN = torch.nn.MSELoss().to(gpu_ids[0])\ncriterion_cycle = torch.nn.L1Loss()\ncriterion_identity = torch.nn.L1Loss()\n\n\ninput_shape = (channels, img_height, img_width)\n\n# Initialize generator and discriminator\ngpu_ids = [0, 1]\nG_AB = init_net(GeneratorResNet(input_shape, n_residual_blocks), gpu_ids=gpu_ids)\nG_BA = init_net(GeneratorResNet(input_shape, n_residual_blocks), gpu_ids=gpu_ids)\nD_A = init_net(Discriminator(input_shape), gpu_ids=gpu_ids)\nD_B = init_net(Discriminator(input_shape), gpu_ids=gpu_ids)\nG_AB = torch.compile(G_AB, mode=\"reduce-overhead\")\nG_BA = torch.compile(G_BA, mode=\"reduce-overhead\")\nD_A = torch.compile(D_A, mode=\"reduce-overhead\")\nD_B= torch.compile(D_B, mode=\"reduce-overhead\")\n\n\noptimizer_G = torch.optim.Adam(\n    itertools.chain(G_AB.parameters(), G_BA.parameters()), lr=lr, betas=(b1, b2)\n)\noptimizer_D_A = torch.optim.Adam(D_A.parameters(), lr=lr, betas=(b1, b2))\noptimizer_D_B = torch.optim.Adam(D_B.parameters(), lr=lr, betas=(b1, b2))\n\nlr_scheduler_G = torch.optim.lr_scheduler.LambdaLR(\n    optimizer_G, lr_lambda=LambdaLR(n_epochs, epoch, decay_epoch).step\n)\nlr_scheduler_D_A = torch.optim.lr_scheduler.LambdaLR(\n    optimizer_D_A, lr_lambda=LambdaLR(n_epochs, epoch, decay_epoch).step\n)\nlr_scheduler_D_B = torch.optim.lr_scheduler.LambdaLR(\n    optimizer_D_B, lr_lambda=LambdaLR(n_epochs, epoch, decay_epoch).step\n)\n\nfake_A_buffer = ReplayBuffer()\nfake_B_buffer = ReplayBuffer()\n\n\n# ----------\n#  Training\n# ----------\n\nprev_time = time.time()\nfor epoch in range(epoch, n_epochs):\n    for i, batch in enumerate(dataloader):\n\n        # Set model input\n        real_A = Variable(batch[\"A\"].cuda())\n        real_B = Variable(batch[\"B\"].cuda())\n\n        # Adversarial ground truths\n        valid = Variable(torch.ones((real_A.size(0), *D_A.module.output_shape)), requires_grad=False).cuda()\n        fake = Variable(torch.zeros((real_A.size(0), *D_A.module.output_shape)), requires_grad=False).cuda()\n\n        # ------------------\n        #  Train Generators\n        # ------------------\n\n        G_AB.train()\n        G_BA.train()\n\n        optimizer_G.zero_grad()\n        with torch.cuda.amp.autocast():\n            # Identity loss\n            loss_id_A = criterion_identity(G_BA(real_A), real_A)\n            loss_id_B = criterion_identity(G_AB(real_B), real_B)\n\n            loss_identity = (loss_id_A + loss_id_B) / 2\n\n            # GAN loss\n            fake_B = G_AB(real_A)\n            loss_GAN_AB = criterion_GAN(D_B(fake_B), valid)\n            fake_A = G_BA(real_B)\n            loss_GAN_BA = criterion_GAN(D_A(fake_A), valid)\n\n            loss_GAN = (loss_GAN_AB + loss_GAN_BA) / 2\n\n            # Cycle loss\n            recov_A = G_BA(fake_B)\n            loss_cycle_A = criterion_cycle(recov_A, real_A)\n            recov_B = G_AB(fake_A)\n            loss_cycle_B = criterion_cycle(recov_B, real_B)\n\n            loss_cycle = (loss_cycle_A + loss_cycle_B) / 2\n\n            # Total loss\n            loss_G = loss_GAN + lambda_cyc * loss_cycle + lambda_id * loss_identity\n\n        loss_G.backward()\n        optimizer_G.step()\n\n        # -----------------------\n        #  Train Discriminator A\n        # -----------------------\n\n        optimizer_D_A.zero_grad()\n        with torch.cuda.amp.autocast():\n\n            # Real loss\n            loss_real = criterion_GAN(D_A(real_A), valid)\n            # Fake loss (on batch of previously generated samples)\n            fake_A_ = fake_A_buffer.push_and_pop(fake_A)\n            loss_fake = criterion_GAN(D_A(fake_A_.detach()), fake)\n            # Total loss\n            loss_D_A = (loss_real + loss_fake) / 2\n\n        loss_D_A.backward()\n        optimizer_D_A.step()\n\n        # -----------------------\n        #  Train Discriminator B\n        # -----------------------\n\n        optimizer_D_B.zero_grad()\n        with torch.cuda.amp.autocast():\n\n            # Real loss\n            loss_real = criterion_GAN(D_B(real_B), valid)\n            # Fake loss (on batch of previously generated samples)\n            fake_B_ = fake_B_buffer.push_and_pop(fake_B)\n            loss_fake = criterion_GAN(D_B(fake_B_.detach()), fake)\n            # Total loss\n            loss_D_B = (loss_real + loss_fake) / 2\n\n        loss_D_B.backward()\n        optimizer_D_B.step()\n\n        loss_D = (loss_D_A + loss_D_B) / 2\n\n        # --------------\n        #  Log Progress\n        # --------------\n\n        # Determine approximate time left\n        batches_done = epoch * len(dataloader) + i\n        batches_left = n_epochs * len(dataloader) - batches_done\n        time_left = datetime.timedelta(seconds=batches_left * (time.time() - prev_time))\n        prev_time = time.time()\n\n        # Print log\n        sys.stdout.write(\n            \"\\r[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f, adv: %f, cycle: %f, identity: %f] ETA: %s\"\n            % (\n                epoch,\n                n_epochs,\n                i,\n                len(dataloader),\n                loss_D.item(),\n                loss_G.item(),\n                loss_GAN.item(),\n                loss_cycle.item(),\n                loss_identity.item(),\n                time_left,\n            )\n        )\n\n    if epoch % 2 == 0 and epoch != 0:\n        sample_images(batches_done, print_image=True)\n\n    # Update learning rates\n    lr_scheduler_G.step()\n    lr_scheduler_D_A.step()\n    lr_scheduler_D_B.step()\n\n    if checkpoint_interval != -1 and epoch % checkpoint_interval == 0:\n        # Save model checkpoints\n        torch.save(G_AB.state_dict(), \"saved_models/G_AB_%d.pth\" % epoch)\n        torch.save(G_BA.state_dict(), \"saved_models/G_BA_%d.pth\" % epoch)\n        torch.save(D_A.state_dict(), \"saved_models/D_A_%d.pth\" % epoch)\n        torch.save(D_B.state_dict(), \"saved_models/D_B_%d.pth\" % epoch)","metadata":{"execution":{"iopub.status.busy":"2023-08-10T09:18:43.012682Z","iopub.execute_input":"2023-08-10T09:18:43.015949Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"os.makedirs(\"/kaggle/working/images\", exist_ok=True)\nG_BA.eval()\nwith torch.no_grad():\n    for i, img in enumerate(val_dataloader):\n        img = Variable(img[\"B\"].type(Tensor))\n        img = G_BA(img)\n        save_image(img, (\"/kaggle/working/images/\" + str(i) + \".jpg\"), normalize=True)\n\nshutil.make_archive(\"/kaggle/working/images\", 'zip', \"/kaggle/working/images\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}