{"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":"# init","metadata":{}},{"cell_type":"code","source":"from __future__ import print_function\n%matplotlib inline\nimport argparse\nimport os\nimport random\nimport torch\nimport torch.nn as nn\nimport torch.nn.parallel\nimport torch.backends.cudnn as cudnn\nimport torch.optim as optim\nimport torch.utils.data\n# import torchvision.datasets as dset\nimport torchvision.transforms as transforms\nimport torchvision.utils as vutils\nimport itertools\nimport numpy as np\nimport matplotlib.pyplot as plt\nfrom PIL import Image\nimport matplotlib.animation as animation\nfrom IPython.display import HTML\n\n# Set random seed for reproducibility\nmanualSeed = 999 #= random.randint(1, 10000) # use if you want new results\nprint(\"Seed: \", manualSeed)\nrandom.seed(manualSeed)\ntorch.manual_seed(manualSeed)","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2023-06-24T05:01:35.274839Z","iopub.execute_input":"2023-06-24T05:01:35.275225Z","iopub.status.idle":"2023-06-24T05:01:35.28947Z","shell.execute_reply.started":"2023-06-24T05:01:35.275197Z","shell.execute_reply":"2023-06-24T05:01:35.288307Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Training arguments","metadata":{}},{"cell_type":"code","source":"class Arguments():\n    # main parameters\n\n    # Root directory for dataset\n    dataroot_A = \"/kaggle/input/gan-getting-started/photo_jpg\"\n    dataroot_B = \"/kaggle/input/gan-getting-started/monet_jpg\"\n    \n    dataroot_test_A = dataroot_A #\"/kaggle/input/gan-getting-started/test_photo\"\n    \n    # Number of workers for dataloader\n    workers = 4\n\n    # Batch size during training\n    batch_size = 4\n    # Spatial size of training images. All images will be resized to this\n    #   size using a transformer.\n    image_size = 256\n\n    # Number of channels in the group A\n    input_nc = 3\n    \n    # Number of channels in the group B\n    output_nc = 3\n\n    # number of filters in the last conv layer of generator\n    ngf = 64\n\n    # number of filters in the first conv layer of discriminator\n    ndf = 64\n\n    # Number of training epochs\n    n_epochs = 1\n\n    # Learning rate for optimizers\n    lr = 0.0002\n    \n    # identity_loss coeff.\n    lambda_identity = .5\n    lambda_A = 10.0\n    lambda_B = 10.0\n\n    # Beta1 hyperparam for Adam optimizers\n    beta1 = 0.5\n\n    # Number of GPUs available. Use 0 for CPU mode.\n    ngpu = 1\n    \n    gan_mode = 'lsgan' # 'vanilla', 'wgangp'\n    \n    pool_size = 16\n\n    device = torch.device(\"cuda:0\" if (torch.cuda.is_available() and ngpu > 0) else \"cpu\")\n    print(device)\n    \n    print_freq = 1\n    save_freq = 1\n    save_dir = \"/kaggle/working/models/\"\n    \nopt = Arguments()\nif not os.path.exists(opt.save_dir):\n    os.makedirs(opt.save_dir)","metadata":{"execution":{"iopub.status.busy":"2023-06-24T05:01:35.295125Z","iopub.execute_input":"2023-06-24T05:01:35.295387Z","iopub.status.idle":"2023-06-24T05:01:35.30415Z","shell.execute_reply.started":"2023-06-24T05:01:35.295364Z","shell.execute_reply":"2023-06-24T05:01:35.30307Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Unpaired training dataset definition\nAdded modified torchvision.ImageFolder","metadata":{}},{"cell_type":"code","source":"\"\"\"A modified image folder class\n\nWe modify the official PyTorch image folder (https://github.com/pytorch/vision/blob/master/torchvision/datasets/folder.py)\nso that this class can load images from both current directory and its subdirectories.\n\"\"\"\n\nIMG_EXTENSIONS = [\n    '.jpg', '.JPG', '.jpeg', '.JPEG',\n    '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP',\n    '.tif', '.TIF', '.tiff', '.TIFF',\n]\n\n\ndef is_image_file(filename):\n    return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)\n\n\ndef make_dataset(dir, max_dataset_size=float(\"inf\")):\n    images = []\n    assert os.path.isdir(dir), '%s is not a valid directory' % dir\n\n    for root, _, fnames in sorted(os.walk(dir)):\n        for fname in fnames:\n            if is_image_file(fname):\n                path = os.path.join(root, fname)\n                images.append(path)\n    return images[:min(max_dataset_size, len(images))]\n\n\ndef default_loader(path):\n    return Image.open(path).convert('RGB')\n\n\nclass ImageFolder(torch.utils.data.Dataset):\n\n    def __init__(self, root, transform=None, return_paths=False,\n                 loader=default_loader):\n        imgs = make_dataset(root)\n        if len(imgs) == 0:\n            raise(RuntimeError(\"Found 0 images in: \" + root + \"\\n\"\n                               \"Supported image extensions are: \" + \",\".join(IMG_EXTENSIONS)))\n\n        self.root = root\n        self.imgs = imgs\n        self.transform = transform\n        self.return_paths = return_paths\n        self.loader = loader\n\n    def __getitem__(self, index):\n        path = self.imgs[index]\n        img = self.loader(path)\n        if self.transform is not None:\n            img = self.transform(img)\n        \n        if self.return_paths:\n            return img, path\n        else:\n            return img\n\n    def __len__(self):\n        return len(self.imgs)\n","metadata":{"execution":{"iopub.status.busy":"2023-06-24T05:01:35.310087Z","iopub.execute_input":"2023-06-24T05:01:35.310621Z","iopub.status.idle":"2023-06-24T05:01:35.324512Z","shell.execute_reply.started":"2023-06-24T05:01:35.310596Z","shell.execute_reply":"2023-06-24T05:01:35.323567Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# create a dataset to return pairs of images from both 2 styles\nclass UnalignedDataset(torch.utils.data.Dataset):\n    def __init__(self, opt):\n        self.path_A = opt.dataroot_A\n        self.path_B = opt.dataroot_B # B: monet\n        \n        self.img_size = opt.image_size\n        \n        transform = transforms.Compose([\n            transforms.Resize(opt.image_size),\n            transforms.CenterCrop(opt.image_size),\n            transforms.ToTensor(),\n            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),\n        ])\n        \n        self.datasetA = ImageFolder(root = self.path_A, transform = transform)\n        self.datasetB = ImageFolder(root = self.path_B, transform = transform)\n        \n        self.size_A = len(self.datasetA)\n        self.size_B = len(self.datasetB)\n\n    def __len__(self):\n        return max(self.size_A, self.size_B)\n\n    def augmentation(self, image):\n        p_rotate = torch.rand(1, dtype=torch.float32)\n        p_spatial = torch.rand(1, dtype=torch.float32)\n        p_crop = torch.rand(1, dtype=torch.float32)\n        transform_list = []\n\n        if p_crop > .5:\n            transform_list.append(transforms.Resize([286, 286], Image.BICUBIC, antialias=True))\n            transform_list.append(transforms.RandomCrop([self.img_size, self.img_size]))\n\n            if p_crop > .9:\n                transform_list.append(transforms.Resize([300, 300], Image.BICUBIC,  antialias=True))\n                transform_list.append(transforms.RandomCrop([self.img_size, self.img_size]))\n        \n        if p_rotate > .9:\n            transform_list.append(transforms.RandomRotation(270))\n        elif p_rotate > .7:\n            transform_list.append(transforms.RandomRotation(180))\n        elif p_rotate > .5:\n            transform_list.append(transforms.RandomRotation(90))\n            \n        # mirroring\n        if p_spatial > .6:\n            transform_list.append(transforms.RandomHorizontalFlip(1.0))\n            transform_list.append(transforms.RandomVerticalFlip(1.0))\n        \n        img_transform = transforms.Compose(transform_list)\n        image = img_transform(image)\n        if p_spatial > .9:\n            image = torch.transpose(image, 1, 2)\n\n        return image\n\n    def __getitem__(self, index):\n        A_tensor = self.datasetA[index % self.size_A]\n        \n        index_B = random.randint(0, self.size_B - 1)\n        B_tensor = self.datasetB[index_B]\n        \n        # do augmentation for monet image\n        B_tensor = self.augmentation(B_tensor)\n        return {'A': A_tensor, 'B': B_tensor}\n\nunpaired_dataset = UnalignedDataset(opt)\ndataloader = torch.utils.data.DataLoader(unpaired_dataset, batch_size=opt.batch_size, shuffle=True, num_workers=opt.workers)","metadata":{"execution":{"iopub.status.busy":"2023-06-24T05:01:35.328583Z","iopub.execute_input":"2023-06-24T05:01:35.328924Z","iopub.status.idle":"2023-06-24T05:01:36.424753Z","shell.execute_reply.started":"2023-06-24T05:01:35.328898Z","shell.execute_reply":"2023-06-24T05:01:36.423796Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Generator(nn.Module):\n    def __init__(self, input_nc, output_nc, ngf):\n        \"\"\"\n        Parameters:\n            input_nc (int) -- the number of channels in input images\n            output_nc (int) -- the number of channels in output images\n            ngf (int) -- the number of filters in the last conv layer\n        Returns a generator\n        \"\"\"\n        super(Generator, self).__init__()\n        self.main = nn.Sequential(\n            # Down sampling\n            nn.Conv2d(input_nc, ngf, 4, 2, 1, bias=False),\n#             nn.BatchNorm2d(ngf),\n            nn.LeakyReLU(0.2, inplace=True),\n            # state size.  ngf x (image_size/2) x (image_size/2)\n            \n            nn.Conv2d(ngf, ngf*2, 4, 2, 1, bias=False),\n#             nn.BatchNorm2d(ngf * 2),\n            nn.LeakyReLU(0.2, inplace=True),\n            # state size. ngf*2 x (image_size/4) x (image_size/4)\n            \n            nn.Conv2d(ngf*2, ngf*4, 4, 2, 1, bias=False),\n#             nn.BatchNorm2d(ngf * 4),\n            nn.LeakyReLU(0.2, inplace=True),\n            # state size. ngf*4 x (image_size/8) x (image_size/8)\n            \n            nn.Conv2d(ngf*4, ngf*8, 4, 2, 1, bias=False),\n            nn.LeakyReLU(0.2, inplace=True),\n            # state size. ngf*8 x (image_size/16) x (image_size/16)\n            \n            # ---\n            # Up sampling\n            nn.ConvTranspose2d(ngf * 8, ngf * 4, 4, 2, 1, bias=False),\n            nn.BatchNorm2d(ngf * 4),\n            nn.ReLU(True),\n            \n            nn.ConvTranspose2d(ngf * 4, ngf * 2, 4, 2, 1, bias=False),\n            nn.BatchNorm2d(ngf * 2),\n            nn.ReLU(True),\n            # state size. ngf*2 x (image_size/4) x (image_size/4)\n            \n            nn.ConvTranspose2d(ngf * 2, ngf, 4, 2, 1, bias=False),\n            nn.BatchNorm2d(ngf),\n            nn.ReLU(True),\n            # state size. (ngf) x image_size/2 x image_size/2\n            \n            nn.ConvTranspose2d(ngf, output_nc, 4, 2, 1, bias=False),\n            nn.Tanh()\n            # state size. (output_nc) x image_size x image_size\n        )\n\n    def forward(self, x):\n        return self.main(x)\n    ","metadata":{"execution":{"iopub.status.busy":"2023-06-24T05:01:36.428608Z","iopub.execute_input":"2023-06-24T05:01:36.428907Z","iopub.status.idle":"2023-06-24T05:01:36.441123Z","shell.execute_reply.started":"2023-06-24T05:01:36.428882Z","shell.execute_reply":"2023-06-24T05:01:36.440164Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Discriminator Code    \nclass Discriminator(nn.Module):\n    def __init__(self, input_nc, ndf):\n        \"\"\"\n        Parameters:\n            input_nc (int)  -- the number of channels in input images\n            ndf (int)       -- the number of filters in the last conv layer\n        \"\"\"\n        super(Discriminator, self).__init__()\n        self.main = nn.Sequential(\n            # input is (input_nc) x image_size x image_size\n            nn.Conv2d(input_nc, ndf, 4, 2, 1, bias=False),\n            nn.BatchNorm2d(ndf),\n            nn.LeakyReLU(0.2, inplace=True),\n            # state size. (ndf) x image_size/2 x image_size/2\n            \n            nn.Conv2d(ndf, ndf * 2, 4, 2, 1, bias=False),\n            nn.BatchNorm2d(ndf * 2),\n            nn.LeakyReLU(0.2, inplace=True),\n            # state size. (ndf*2) x image_size/4 x image_size/4\n            \n            nn.Conv2d(ndf * 2, ndf * 4, 4, 2, 1, bias=False),\n            nn.BatchNorm2d(ndf * 4),\n            nn.LeakyReLU(0.2, inplace=True),\n            # state size. (ndf*4) x image_size/8 x image_size/8\n            \n            nn.Conv2d(ndf * 4, ndf * 8, 4, 2, 1, bias=False),\n            nn.BatchNorm2d(ndf * 8),\n            nn.LeakyReLU(0.2, inplace=True),\n\n            nn.Conv2d(ndf * 8, 1, 4, 1, 0, bias=False),\n            nn.Sigmoid()\n        )\n\n    def forward(self, input):\n        return self.main(input)","metadata":{"execution":{"iopub.status.busy":"2023-06-24T05:01:36.442665Z","iopub.execute_input":"2023-06-24T05:01:36.443237Z","iopub.status.idle":"2023-06-24T05:01:36.456043Z","shell.execute_reply.started":"2023-06-24T05:01:36.443205Z","shell.execute_reply":"2023-06-24T05:01:36.454993Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# custom weights initialization called on netG and netD\ndef weights_init(m):\n    classname = m.__class__.__name__\n    if classname.find('Conv') != -1:\n        nn.init.normal_(m.weight.data, 0.0, 0.02)\n    elif classname.find('BatchNorm') != -1:\n        nn.init.normal_(m.weight.data, 1.0, 0.02)\n        nn.init.constant_(m.bias.data, 0)","metadata":{"execution":{"iopub.status.busy":"2023-06-24T05:01:36.458788Z","iopub.execute_input":"2023-06-24T05:01:36.459117Z","iopub.status.idle":"2023-06-24T05:01:36.468278Z","shell.execute_reply.started":"2023-06-24T05:01:36.459083Z","shell.execute_reply":"2023-06-24T05:01:36.467312Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class GANLoss(nn.Module):\n    \"\"\"Define different GAN objectives.\n\n    The GANLoss class abstracts away the need to create the target label tensor\n    that has the same size as the input.\n    \"\"\"\n\n    def __init__(self, gan_mode, target_real_label=1.0, target_fake_label=0.0):\n        \"\"\" Initialize the GANLoss class.\n\n        Parameters:\n            gan_mode (str) - - the type of GAN objective. It currently supports vanilla, lsgan, and wgangp.\n            target_real_label (bool) - - label for a real image\n            target_fake_label (bool) - - label of a fake image\n\n        Note: Do not use sigmoid as the last layer of Discriminator.\n        LSGAN needs no sigmoid. vanilla GANs will handle it with BCEWithLogitsLoss.\n        \"\"\"\n        super(GANLoss, self).__init__()\n        self.register_buffer('real_label', torch.tensor(target_real_label))\n        self.register_buffer('fake_label', torch.tensor(target_fake_label))\n        self.gan_mode = gan_mode\n        if gan_mode == 'lsgan':\n            self.loss = nn.MSELoss()\n        elif gan_mode == 'vanilla':\n            self.loss = nn.BCEWithLogitsLoss()\n        elif gan_mode in ['wgangp']:\n            self.loss = None\n        else:\n            raise NotImplementedError('gan mode %s not implemented' % gan_mode)\n\n    def get_target_tensor(self, prediction, target_is_real):\n        \"\"\"Create label tensors with the same size as the input.\n\n        Parameters:\n            prediction (tensor) - - tpyically the prediction from a discriminator\n            target_is_real (bool) - - if the ground truth label is for real images or fake images\n\n        Returns:\n            A label tensor filled with ground truth label, and with the size of the input\n        \"\"\"\n\n        if target_is_real:\n            target_tensor = self.real_label\n        else:\n            target_tensor = self.fake_label\n        return target_tensor.expand_as(prediction)\n\n    def __call__(self, prediction, target_is_real):\n        \"\"\"Calculate loss given Discriminator's output and ground truth labels.\n\n        Parameters:\n            prediction (tensor) - - tpyically the prediction output from a discriminator\n            target_is_real (bool) - - if the ground truth label is for real images or fake images\n\n        Returns:\n            the calculated loss.\n        \"\"\"\n        if self.gan_mode in ['lsgan', 'vanilla']:\n            target_tensor = self.get_target_tensor(prediction, target_is_real)\n            loss = self.loss(prediction, target_tensor)\n        elif self.gan_mode == 'wgangp':\n            if target_is_real:\n                loss = -prediction.mean()\n            else:\n                loss = prediction.mean()\n        return loss","metadata":{"execution":{"iopub.status.busy":"2023-06-24T05:01:36.470093Z","iopub.execute_input":"2023-06-24T05:01:36.470421Z","iopub.status.idle":"2023-06-24T05:01:36.483078Z","shell.execute_reply.started":"2023-06-24T05:01:36.470391Z","shell.execute_reply":"2023-06-24T05:01:36.481797Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class ImagePool():\n    \"\"\"This class implements an image buffer that stores previously generated images.\n\n    This buffer enables us to update discriminators using a history of generated images\n    rather than the ones produced by the latest generators.\n    \"\"\"\n\n    def __init__(self, pool_size):\n        \"\"\"Initialize the ImagePool class\n\n        Parameters:\n            pool_size (int) -- the size of image buffer, if pool_size=0, no buffer will be created\n        \"\"\"\n        self.pool_size = pool_size\n        if self.pool_size > 0:  # create an empty pool\n            self.num_imgs = 0\n            self.images = []\n\n    def query(self, images):\n        \"\"\"Return an image from the pool.\n\n        Parameters:\n            images: the latest generated images from the generator\n\n        Returns images from the buffer.\n\n        By 50/100, the buffer will return input images.\n        By 50/100, the buffer will return images previously stored in the buffer,\n        and insert the current images to the buffer.\n        \"\"\"\n        if self.pool_size == 0:  # if the buffer size is 0, do nothing\n            return images\n        return_images = []\n        for image in images:\n            image = torch.unsqueeze(image.data, 0)\n            if self.num_imgs < self.pool_size:   # if the buffer is not full; keep inserting current images to the buffer\n                self.num_imgs = self.num_imgs + 1\n                self.images.append(image)\n                return_images.append(image)\n            else:\n                p = random.uniform(0, 1)\n                if p > 0.5:  # by 50% chance, the buffer will return a previously stored image, and insert the current image into the buffer\n                    random_id = random.randint(0, self.pool_size - 1)  # randint is inclusive\n                    tmp = self.images[random_id].clone()\n                    self.images[random_id] = image\n                    return_images.append(tmp)\n                else:       # by another 50% chance, the buffer will return the current image\n                    return_images.append(image)\n        return_images = torch.cat(return_images, 0)   # collect all the images and return\n        return return_images","metadata":{"execution":{"iopub.status.busy":"2023-06-24T05:01:36.485809Z","iopub.execute_input":"2023-06-24T05:01:36.486519Z","iopub.status.idle":"2023-06-24T05:01:36.497078Z","shell.execute_reply.started":"2023-06-24T05:01:36.486487Z","shell.execute_reply":"2023-06-24T05:01:36.495965Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CycleGANModel():\n    def __init__(self, opt):\n        # specify the training losses you want to print out.\n        \n        self.opt = opt\n        self.loss_names = ['D_A', 'G_A', 'cycle_A', 'idt_A', 'D_B', 'G_B', 'cycle_B', 'idt_B']\n        \n        # specify the images you want to save/display.\n        visual_names_A = ['real_A', 'fake_B', 'rec_A', 'idt_B']\n        visual_names_B = ['real_B', 'fake_A', 'rec_B', 'idt_A']\n        \n        self.visual_names = visual_names_A + visual_names_B  # combine visualizations for A and B\n        \n        self.model_names = ['G_A', 'G_B', 'D_A', 'D_B']\n        self.save_dir = opt.save_dir\n        \n        self.netG_A = Generator(opt.input_nc, opt.output_nc, opt.ngf).to(opt.device)\n        self.netG_B = Generator(opt.output_nc, opt.input_nc, opt.ngf).to(opt.device)\n\n        self.netD_A = Discriminator(opt.output_nc, opt.ndf).to(opt.device)\n        self.netD_B = Discriminator(opt.input_nc, opt.ndf).to(opt.device)\n        \n        self.netG_A.apply(weights_init)\n        self.netG_B.apply(weights_init)\n        self.netD_A.apply(weights_init)\n        self.netD_B.apply(weights_init)\n        \n        self.fake_A_pool = ImagePool(opt.pool_size)  # create image buffer to store previously generated images\n        self.fake_B_pool = ImagePool(opt.pool_size)  # create image buffer to store previously generated images\n            \n        # define loss functions\n        self.criterionGAN = GANLoss(opt.gan_mode).to(opt.device)\n        self.criterionCycle = torch.nn.L1Loss()\n        self.criterionIdt = torch.nn.L1Loss()\n        \n        # initialize optimizers; schedulers will be automatically created by function <BaseModel.setup>.\n        self.optimizers = []\n        self.optimizer_G = torch.optim.Adam(itertools.chain(self.netG_A.parameters(), self.netG_B.parameters()), lr=opt.lr, betas=(opt.beta1, 0.999))\n        self.optimizer_D = torch.optim.Adam(itertools.chain(self.netD_A.parameters(), self.netD_B.parameters()), lr=opt.lr, betas=(opt.beta1, 0.999))\n        self.optimizers.append(self.optimizer_G)\n        self.optimizers.append(self.optimizer_D)\n        \n    def set_requires_grad(self, nets, requires_grad=False):\n        if not isinstance(nets, list):\n            nets = [nets]\n        for net in nets:\n            if net is not None:\n                for param in net.parameters():\n                    param.requires_grad = requires_grad\n                    \n    def set_input(self, input):\n        \"\"\"Unpack input data from the dataloader\"\"\"\n        self.real_A = input['A'].to(self.opt.device)\n        self.real_B = input['B'].to(self.opt.device)\n        return\n        \n    def forward(self):\n        \"\"\"Run forward pass\"\"\"\n        self.fake_B = self.netG_A(self.real_A)  # G_A(A)\n        self.rec_A = self.netG_B(self.fake_B)   # G_B(G_A(A))\n        self.fake_A = self.netG_B(self.real_B)  # G_B(B)\n        self.rec_B = self.netG_A(self.fake_A)   # G_A(G_B(B))\n        return\n    \n    def backward_D_basic(self, netD, real, fake):\n        # Real\n        pred_real = netD(real)\n        loss_D_real = self.criterionGAN(pred_real, True)\n        # Fake\n        pred_fake = netD(fake.detach())\n        loss_D_fake = self.criterionGAN(pred_fake, False)\n        # Combined loss and calculate gradients\n        loss_D = (loss_D_real + loss_D_fake) * 0.5\n        loss_D.backward()\n        return loss_D\n    \n    def backward_D_A(self):\n        \"\"\"Calculate GAN loss for discriminator D_A\"\"\"\n        fake_B = self.fake_B_pool.query(self.fake_B)\n        self.loss_D_A = self.backward_D_basic(self.netD_A, self.real_B, fake_B)\n        return\n    \n    def backward_D_B(self):\n        \"\"\"Calculate GAN loss for discriminator D_B\"\"\"\n        fake_A = self.fake_A_pool.query(self.fake_A)\n        self.loss_D_B = self.backward_D_basic(self.netD_B, self.real_A, fake_A)\n        return\n    \n    def backward_G(self):\n        \"\"\"Calculate the loss for generators G_A and G_B\"\"\"\n        lambda_idt = self.opt.lambda_identity\n        lambda_A = self.opt.lambda_A\n        lambda_B = self.opt.lambda_B\n        # Identity loss\n        if lambda_idt > 0:\n            # G_A should be identity if real_B is fed: ||G_A(B) - B||\n            self.idt_A = self.netG_A(self.real_B)\n            self.loss_idt_A = self.criterionIdt(self.idt_A, self.real_B) * lambda_B * lambda_idt\n            # G_B should be identity if real_A is fed: ||G_B(A) - A||\n            self.idt_B = self.netG_B(self.real_A)\n            self.loss_idt_B = self.criterionIdt(self.idt_B, self.real_A) * lambda_A * lambda_idt\n        else:\n            self.loss_idt_A = 0\n            self.loss_idt_B = 0\n\n        # GAN loss D_A(G_A(A))\n        self.loss_G_A = self.criterionGAN(self.netD_A(self.fake_B), True)\n        # GAN loss D_B(G_B(B))\n        self.loss_G_B = self.criterionGAN(self.netD_B(self.fake_A), True)\n        # Forward cycle loss || G_B(G_A(A)) - A||\n        self.loss_cycle_A = self.criterionCycle(self.rec_A, self.real_A) * lambda_A\n        # Backward cycle loss || G_A(G_B(B)) - B||\n        self.loss_cycle_B = self.criterionCycle(self.rec_B, self.real_B) * lambda_B\n        # combined loss and calculate gradients\n        self.loss_G = self.loss_G_A + self.loss_G_B + self.loss_cycle_A + self.loss_cycle_B + self.loss_idt_A + self.loss_idt_B\n        self.loss_G.backward()\n        return\n    \n    def optimize_parameters(self):\n        \"\"\"Calculate losses, gradients, and update network weights; called in every training iteration\"\"\"\n        # forward\n        self.forward()      # compute fake images and reconstruction images.\n        # G_A and G_B\n        self.set_requires_grad([self.netD_A, self.netD_B], False)  # Ds require no gradients when optimizing Gs\n        self.optimizer_G.zero_grad()  # set G_A and G_B's gradients to zero\n        self.backward_G()             # calculate gradients for G_A and G_B\n        self.optimizer_G.step()       # update G_A and G_B's weights\n        # D_A and D_B\n        self.set_requires_grad([self.netD_A, self.netD_B], True)\n        self.optimizer_D.zero_grad()   # set D_A and D_B's gradients to zero\n        self.backward_D_A()      # calculate gradients for D_A\n        self.backward_D_B()      # calculate graidents for D_B\n        self.optimizer_D.step()  # update D_A and D_B's weights\n\n    def get_current_losses(self):\n        losses_dict = {}\n        for name in self.loss_names:\n            losses_dict[name] = float(getattr(self, 'loss_' + name))\n        return losses_dict\n    \n    def save_networks(self, epoch):\n        for name in self.model_names:\n            if isinstance(name, str):\n                save_filename = f'{epoch}_net_{name}.pth'\n                save_path = os.path.join(self.save_dir, save_filename)\n                net = getattr(self, 'net' + name)\n\n                if torch.cuda.is_available():\n                    torch.save(net.state_dict(), save_path)\n                    # net.cuda(0)\n                else:\n                    torch.save(net.cpu().state_dict(), save_path)\n        return\n\n    \n    def load_networks(self, epoch):\n        for name in self.model_names:\n            if isinstance(name, str):\n                load_filename = f'{epoch}_net_{name}.pth'\n                load_path = os.path.join(self.save_dir, load_filename)\n                net = getattr(self, 'net' + name)\n                \n                net.load_state_dict(torch.load(load_path, map_location=self.opt.device))\n        return\n    ","metadata":{"execution":{"iopub.status.busy":"2023-06-24T05:01:36.500509Z","iopub.execute_input":"2023-06-24T05:01:36.501956Z","iopub.status.idle":"2023-06-24T05:01:36.508343Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# # initialize model\n# model = CycleGANModel(opt)\n# current_epoch = 0\n\n# # model.load_networks(epoch = current_epoch)\n# # main loop\n# img_list =[]\n# G_losses = []\n# D_losses = []\n\n# iter_print_freq = 100\n# iter_save_freq = 1000\n\n# for epoch in range(current_epoch, current_epoch + opt.n_epochs):\n#     print(f\"=============epoch {epoch} =============\")\n#     for i, data in enumerate(dataloader, start=1):\n                \n#         model.set_input(data)         # unpack data from dataset and apply preprocessing\n#         model.optimize_parameters()   # calculate loss functions, get gradients, update network weights\n        \n#         # Save Losses for plotting later\n#         G_losses.append(model.loss_G_A.item() + model.loss_G_B.item())\n#         D_losses.append(model.loss_D_A.item() + model.loss_D_B.item())\n#         if i % iter_print_freq == 0:\n#             losses = model.get_current_losses()\n#             print(f\"iter_{i}, G: {model.loss_G:.3f}, GA: {model.loss_G_A}, GB: {model.loss_G_B}, DA: {model.loss_D_A}, DB: {model.loss_D_B}, CycleA: {model.loss_cycle_A}, CycleB: {model.loss_cycle_B}\")\n#         if i % iter_save_freq == 0:\n#             print(f\"saving the latest model, epoch: {epoch + 1}.\")\n#             model.save_networks(epoch + 1)\n        \n#     # progress of the network for the fixed_test_batch\n#     with torch.no_grad():\n#         fake = model.netG_A(fixed_test_batch).detach().cpu()\n#     img_list.append(vutils.make_grid(fake, padding=2, normalize=True))\n        \n#     # save networks\n#     if (epoch + 1) % opt.save_freq == 0 or epoch == opt.n_epochs - 1:\n#         print(f\"saving the latest model, epoch: {epoch + 1}.\")\n#         model.save_networks(epoch + 1)\n        \n#     # plot losses\n#     if (epoch + 1) % opt.print_freq == 0:\n#         losses = model.get_current_losses()\n\n#         print(f\"---{epoch+1}---: {model.loss_G:.3f}\")\n#         for l,v in losses.items():\n#             print(f\"loss {l}: {v:.3f}\")\n        ","metadata":{"execution":{"iopub.status.busy":"2023-06-24T05:01:36.509113Z","iopub.execute_input":"2023-06-24T05:01:36.509354Z","iopub.status.idle":"2023-06-24T05:01:41.405287Z","shell.execute_reply.started":"2023-06-24T05:01:36.509333Z","shell.execute_reply":"2023-06-24T05:01:41.403635Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# evaluation A->B","metadata":{}},{"cell_type":"code","source":"# dataset_test = ImageFolder(\n#     root = opt.dataroot_test_A,\n#     transform = transforms.Compose([\n#         transforms.Resize(opt.image_size),\n#         transforms.CenterCrop(opt.image_size),\n#         transforms.ToTensor(),\n#         transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),\n#     ]))\n\n# dataloader_test = torch.utils.data.DataLoader(dataset_test, batch_size=4, shuffle=False, num_workers=0)\n\n# # # fixed batch used to visualize the evolution of the network over training\n# fixed_test_batch = next(iter(dataloader_test)).to(opt.device)\n\n\n# grid_real = vutils.make_grid(fixed_test_batch.cpu(), padding=2, normalize=True, nrow=2)\n# with torch.no_grad():\n#     fake_batch = model.netG_A(fixed_test_batch).detach().cpu()\n# #     fake_batch = model.netG_A(model.netG_A(real_batch)).detach().cpu() # double the effect\n\n# plt.figure(figsize=(16,8), dpi=128)\n\n# plt.subplot(1,2,1)\n# plt.axis(\"off\")\n# grid_fake = vutils.make_grid(fake_batch, padding=2, normalize=True, nrow=2)\n# plt.imshow(np.transpose(grid_fake,(1,2,0)))\n\n# plt.subplot(1,2,2)\n# plt.axis(\"off\")\n# plt.imshow(np.transpose(grid_real,(1,2,0)))","metadata":{"execution":{"iopub.status.busy":"2023-06-24T05:01:41.406625Z","iopub.status.idle":"2023-06-24T05:01:41.40743Z","shell.execute_reply.started":"2023-06-24T05:01:41.407178Z","shell.execute_reply":"2023-06-24T05:01:41.407203Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# evaluation  B->A","metadata":{}},{"cell_type":"code","source":"# with torch.no_grad():\n#     fake_batch = model.netG_B(fixed_test_batch).detach().cpu()\n# #     fake_batch = model.netG_B(model.netG_B(real_batch)).detach().cpu() # double effect\n\n# plt.figure(figsize=(16,8), dpi=128)\n\n# plt.subplot(1,2,1)\n# plt.axis(\"off\")\n# grid_fake = vutils.make_grid(fake_batch, padding=2, normalize=True, nrow=2)\n# plt.imshow(np.transpose(grid_fake,(1,2,0)))\n\n# plt.subplot(1,2,2)\n# plt.axis(\"off\")\n# plt.imshow(np.transpose(grid_real,(1,2,0)))\n\n","metadata":{"execution":{"iopub.status.busy":"2023-06-24T05:01:41.409463Z","iopub.status.idle":"2023-06-24T05:01:41.410562Z","shell.execute_reply.started":"2023-06-24T05:01:41.410305Z","shell.execute_reply":"2023-06-24T05:01:41.410327Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"","metadata":{}},{"cell_type":"markdown","source":"# do style transfer","metadata":{}},{"cell_type":"code","source":"import shutil\n\nnetG_A = Generator(opt.input_nc, opt.output_nc, opt.ngf)\nnetG_A.load_state_dict(torch.load(os.path.join(opt.save_dir, \"1_net_G_A.pth\"), map_location=str(opt.device)))\nnetG_A = netG_A.to(opt.device)\n\ndataset_test = ImageFolder(\n    root = opt.dataroot_test_A,\n    transform = transforms.Compose([\n        transforms.Resize(opt.image_size),\n        transforms.CenterCrop(opt.image_size),\n        transforms.ToTensor(),\n        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),\n    ]))\n\ndataloader_test = torch.utils.data.DataLoader(dataset_test, batch_size=4, shuffle=False, num_workers=0)\n\nimages = []\ntransform = transforms.ToPILImage()\n\nwith torch.no_grad():\n    for i, data in enumerate(dataloader_test, start=1):\n        data = data.to(opt.device)\n        fake_batch = netG_A(data).detach().cpu()\n        for b in range(fake_batch.shape[0]):\n            img = transform(fake_batch[b])\n            images.append(img)\n\ntmp_dir: str = '../tmp'\n        \nif not os.path.exists(tmp_dir):\n    os.makedirs(tmp_dir)\n\ncnt = 1\nfor img in images:\n    img.save(os.path.join(tmp_dir, f\"{cnt}.jpg\"))\n    cnt = cnt + 1\n\nshutil.make_archive(\"/kaggle/images\", 'zip', tmp_dir)\nshutil.rmtree(tmp_dir)","metadata":{"execution":{"iopub.status.busy":"2023-06-24T05:16:37.46573Z","iopub.execute_input":"2023-06-24T05:16:37.466474Z","iopub.status.idle":"2023-06-24T05:17:33.932134Z","shell.execute_reply.started":"2023-06-24T05:16:37.466434Z","shell.execute_reply":"2023-06-24T05:17:33.931063Z"},"trusted":true},"execution_count":null,"outputs":[]}]}