{"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":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\ni = 0\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        if i%4000 ==0:\n            print(os.path.join(dirname, filename))\n        i +=1\nprint('Done')\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2021-08-18T17:49:16.849886Z","iopub.execute_input":"2021-08-18T17:49:16.850246Z","iopub.status.idle":"2021-08-18T17:49:18.936674Z","shell.execute_reply.started":"2021-08-18T17:49:16.850162Z","shell.execute_reply":"2021-08-18T17:49:18.93344Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport pandas as pd\nimport cv2\nfrom tqdm.auto import tqdm\nfrom torch.utils.data import Dataset,DataLoader\nfrom torchvision.utils import make_grid\nimport math\nfrom PIL import Image","metadata":{"execution":{"iopub.status.busy":"2021-08-18T17:49:18.937506Z","iopub.status.idle":"2021-08-18T17:49:18.937883Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"main_path = \"/kaggle/input/gan-getting-started/\"\nmonet_data  = main_path + \"monet_jpg/\"\nphoto_data = main_path + \"photo_jpg/\"\nmonet_data,photo_data","metadata":{"execution":{"iopub.status.busy":"2021-08-18T17:49:18.938914Z","iopub.status.idle":"2021-08-18T17:49:18.939436Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import torch\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\n\nDEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nTRAIN_DIR = \"/kaggle/input/gan-getting-started\"\nVAL_DIR = \"data/val\"\nBATCH_SIZE = 1\nLEARNING_RATE = 1e-5\nLAMBDA_IDENTITY = 0.0\nLAMBDA_CYCLE = 10\nNUM_WORKERS = 4\nNUM_EPOCHS = 10\nLOAD_MODEL = True\nSAVE_MODEL = True\nCHECKPOINT_GEN_H = \"genh.pth.tar\"\nCHECKPOINT_GEN_Z = \"genz.pth.tar\"\nCHECKPOINT_CRITIC_H = \"critich.pth.tar\"\nCHECKPOINT_CRITIC_Z = \"criticz.pth.tar\"\n\ntransforms = A.Compose(\n    [\n        A.Resize(width=256, height=256),\n        A.HorizontalFlip(p=0.5),\n        A.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], max_pixel_value=255),\n        ToTensorV2(),\n     ],\n    additional_targets={\"image0\": \"image\"},\n)","metadata":{"execution":{"iopub.status.busy":"2021-08-18T17:49:18.940367Z","iopub.status.idle":"2021-08-18T17:49:18.940901Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#utils\nimport random, torch, os, numpy as np\nimport torch.nn as nn\nimport copy\n\ndef save_checkpoint(model, optimizer, filename=\"my_checkpoint.pth.tar\"):\n    print(\"=> Saving checkpoint\")\n    checkpoint = {\n        \"state_dict\": model.state_dict(),\n        \"optimizer\": optimizer.state_dict(),\n    }\n    torch.save(checkpoint, filename)\n\n\ndef load_checkpoint(checkpoint_file, model, optimizer, lr):\n    print(\"=> Loading checkpoint\")\n    checkpoint = torch.load(checkpoint_file, map_location=config.DEVICE)\n    model.load_state_dict(checkpoint[\"state_dict\"])\n    optimizer.load_state_dict(checkpoint[\"optimizer\"])\n\n    for param_group in optimizer.param_groups:\n        param_group[\"lr\"] = lr\n\n\ndef seed_everything(seed=42):\n    os.environ[\"PYTHONHASHSEED\"] = str(seed)\n    random.seed(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.cuda.manual_seed_all(seed)\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = False","metadata":{"execution":{"iopub.status.busy":"2021-08-18T17:49:18.941877Z","iopub.status.idle":"2021-08-18T17:49:18.94251Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def print_details(path):\n    folder_files_count = 0\n    image_shape = {}\n    for i in os.listdir(path):\n        image = cv2.imread(path+i)\n        image_shape['image_shape'] = image.shape\n        image_shape['folder_name'] = path\n        folder_files_count += 1\n    image_shape['items'] = folder_files_count\n    return image_shape\nprint_details(monet_data)\n        ","metadata":{"execution":{"iopub.status.busy":"2021-08-18T17:49:18.943674Z","iopub.status.idle":"2021-08-18T17:49:18.944292Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# print_details(photo_data)","metadata":{"execution":{"iopub.status.busy":"2021-08-18T17:49:18.945397Z","iopub.status.idle":"2021-08-18T17:49:18.946028Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def see_batch(path,n_images,fig_size=(25,25)):\n    plt.figure(figsize = fig_size)\n    w = int(n_images ** .5)\n    h = math.ceil(n_images / w)\n    \n    all_names = os.listdir(path)\n    \n    image_names = all_names[:n_images]\n    for ind, image_name in enumerate(image_names):\n        img = cv2.imread(os.path.join(path, image_name))\n        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) \n        plt.subplot(h, w, ind + 1)\n        plt.imshow(img)\n        plt.axis(\"off\")\n    \n    plt.show()","metadata":{"execution":{"iopub.status.busy":"2021-08-18T17:49:18.947213Z","iopub.status.idle":"2021-08-18T17:49:18.947843Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# see_batch(monet_data,25)","metadata":{"execution":{"iopub.status.busy":"2021-08-18T17:49:18.948912Z","iopub.status.idle":"2021-08-18T17:49:18.949569Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# see_batch(photo_data,25)","metadata":{"execution":{"iopub.status.busy":"2021-08-18T17:49:18.950709Z","iopub.status.idle":"2021-08-18T17:49:18.951287Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Monet2Phots(Dataset):\n    def __init__(self,root_monet,root_picture,transform = None):\n        self.root_monet = root_monet\n        self.root_picture = root_picture\n        self.transform = transform\n        \n        self.monet_images = os.listdir(root_monet)\n        self.picture_images = os.listdir(root_picture)\n        \n        self.length_dataset = max(len(self.monet_images),len(self.picture_images)) - 6500\n        \n        self.monet_len = len(self.monet_images)\n        self.picture_len = len(self.picture_images) - 6500\n        \n    def __len__(self):\n        return self.length_dataset\n    \n    def __getitem__(self,index):\n        monet_img = self.monet_images[index%self.monet_len]\n        picture_img = self.picture_images[index%self.picture_len]\n        \n        monet_path = os.path.join(self.root_monet,monet_img)\n        picture_path = os.path.join(self.root_picture,picture_img)\n        \n        monet_img = np.array(Image.open(monet_path).convert(\"RGB\"))\n        picture_img = np.array(Image.open(picture_path).convert(\"RGB\"))\n        \n        if self.transform:\n            argumentaions = self.transform(image = monet_img,image0 = picture_img)\n            monet_img = argumentaions['image']\n            picture_img = argumentaions['image0']\n        return monet_img,picture_img\n        \n        ","metadata":{"execution":{"iopub.status.busy":"2021-08-18T17:49:18.952357Z","iopub.status.idle":"2021-08-18T17:49:18.952956Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"dataset = Monet2Phots(monet_data,photo_data,transforms)","metadata":{"execution":{"iopub.status.busy":"2021-08-18T17:49:18.953996Z","iopub.status.idle":"2021-08-18T17:49:18.954611Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"dataset.length_dataset","metadata":{"execution":{"iopub.status.busy":"2021-08-18T17:49:18.955824Z","iopub.status.idle":"2021-08-18T17:49:18.956431Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"loader = DataLoader(dataset,\n                    batch_size = BATCH_SIZE,shuffle = True,\n                    num_workers = NUM_WORKERS,pin_memory = True)","metadata":{"execution":{"iopub.status.busy":"2021-08-18T17:49:18.957524Z","iopub.status.idle":"2021-08-18T17:49:18.958128Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class ResidualBlock(nn.Module):\n    def __init__(self, input_channels):\n        super(ResidualBlock, self).__init__()\n        self.conv1 = nn.Conv2d(input_channels, input_channels, kernel_size=3, padding=1, padding_mode='reflect')\n        self.conv2 = nn.Conv2d(input_channels, input_channels, kernel_size=3, padding=1, padding_mode='reflect')\n        self.instancenorm = nn.InstanceNorm2d(input_channels)\n        self.activation = nn.ReLU()\n\n    def forward(self, x):\n        original_x = x.clone()\n        x = self.conv1(x)\n        x = self.instancenorm(x)\n        x = self.activation(x)\n        x = self.conv2(x)\n        x = self.instancenorm(x)\n        return original_x + x","metadata":{"execution":{"iopub.status.busy":"2021-08-18T17:49:19.031946Z","iopub.status.idle":"2021-08-18T17:49:19.032601Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class ContractingBlock(nn.Module):\n    def __init__(self, input_channels, use_bn=True, kernel_size=3, activation='relu'):\n        super(ContractingBlock, self).__init__()\n        self.conv1 = nn.Conv2d(input_channels, input_channels * 2, kernel_size=kernel_size, padding=1, stride=2, padding_mode='reflect')\n        self.activation = nn.ReLU() if activation == 'relu' else nn.LeakyReLU(0.2)\n        if use_bn:\n            self.instancenorm = nn.InstanceNorm2d(input_channels * 2)\n        self.use_bn = use_bn\n\n    def forward(self, x):\n        x = self.conv1(x)\n        if self.use_bn:\n            x = self.instancenorm(x)\n        x = self.activation(x)\n        return x\n\nclass ExpandingBlock(nn.Module):\n    def __init__(self, input_channels, use_bn=True):\n        super(ExpandingBlock, self).__init__()\n        self.conv1 = nn.ConvTranspose2d(input_channels, input_channels // 2, kernel_size=3, stride=2, padding=1, output_padding=1)\n        if use_bn:\n            self.instancenorm = nn.InstanceNorm2d(input_channels // 2)\n        self.use_bn = use_bn\n        self.activation = nn.ReLU()\n\n    def forward(self, x):\n        x = self.conv1(x)\n        if self.use_bn:\n            x = self.instancenorm(x)\n        x = self.activation(x)\n        return x\n\nclass FeatureMapBlock(nn.Module):\n    def __init__(self, input_channels, output_channels):\n        super(FeatureMapBlock, self).__init__()\n        self.conv = nn.Conv2d(input_channels, output_channels, kernel_size=7, padding=3, padding_mode='reflect')\n\n    def forward(self, x):\n        x = self.conv(x)\n        return x","metadata":{"execution":{"iopub.status.busy":"2021-08-18T17:49:19.033879Z","iopub.status.idle":"2021-08-18T17:49:19.034488Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Generator(nn.Module):\n    def __init__(self, input_channels, output_channels, hidden_channels=64):\n        super(Generator, self).__init__()\n        self.upfeature = FeatureMapBlock(input_channels, hidden_channels)\n        self.contract1 = ContractingBlock(hidden_channels)\n        self.contract2 = ContractingBlock(hidden_channels * 2)\n        res_mult = 4\n        self.res0 = ResidualBlock(hidden_channels * res_mult)\n        self.res1 = ResidualBlock(hidden_channels * res_mult)\n        self.res2 = ResidualBlock(hidden_channels * res_mult)\n        self.res3 = ResidualBlock(hidden_channels * res_mult)\n        self.res4 = ResidualBlock(hidden_channels * res_mult)\n        self.res5 = ResidualBlock(hidden_channels * res_mult)\n        self.res6 = ResidualBlock(hidden_channels * res_mult)\n        self.res7 = ResidualBlock(hidden_channels * res_mult)\n        self.res8 = ResidualBlock(hidden_channels * res_mult)\n        self.expand2 = ExpandingBlock(hidden_channels * 4)\n        self.expand3 = ExpandingBlock(hidden_channels * 2)\n        self.downfeature = FeatureMapBlock(hidden_channels, output_channels)\n        self.tanh = torch.nn.Tanh()\n\n    def forward(self, x):\n        x0 = self.upfeature(x)\n        x1 = self.contract1(x0)\n        x2 = self.contract2(x1)\n        x3 = self.res0(x2)\n        x4 = self.res1(x3)\n        x5 = self.res2(x4)\n        x6 = self.res3(x5)\n        x7 = self.res4(x6)\n        x8 = self.res5(x7)\n        x9 = self.res6(x8)\n        x10 = self.res7(x9)\n        x11 = self.res8(x10)\n        x12 = self.expand2(x11)\n        x13 = self.expand3(x12)\n        xn = self.downfeature(x13)\n        return self.tanh(xn)","metadata":{"execution":{"iopub.status.busy":"2021-08-18T17:49:19.035628Z","iopub.status.idle":"2021-08-18T17:49:19.036275Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Discriminator(nn.Module):\n    def __init__(self, input_channels, hidden_channels=64):\n        super(Discriminator, self).__init__()\n        self.upfeature = FeatureMapBlock(input_channels, hidden_channels)\n        self.contract1 = ContractingBlock(hidden_channels, use_bn=False, kernel_size=4, activation='lrelu')\n        self.contract2 = ContractingBlock(hidden_channels * 2, kernel_size=4, activation='lrelu')\n        self.contract3 = ContractingBlock(hidden_channels * 4, kernel_size=4, activation='lrelu')\n        self.final = nn.Conv2d(hidden_channels * 8, 1, kernel_size=1)\n\n    def forward(self, x):\n        x0 = self.upfeature(x)\n        x1 = self.contract1(x0)\n        x2 = self.contract2(x1)\n        x3 = self.contract3(x2)\n        xn = self.final(x3)\n        return xn","metadata":{"execution":{"iopub.status.busy":"2021-08-18T17:49:19.037356Z","iopub.status.idle":"2021-08-18T17:49:19.037976Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import torch.nn.functional as F\n\nadv_criterion = nn.MSELoss() \nrecon_criterion = nn.L1Loss() \n\nn_epochs = 200\ndim_A = 3\ndim_B = 3\ndisplay_step = 200\nlr = 0.0002\nload_shape = 286\ntarget_shape = 256\ndevice = 'cuda'\n","metadata":{"execution":{"iopub.status.busy":"2021-08-18T17:49:19.039142Z","iopub.status.idle":"2021-08-18T17:49:19.039747Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"gen_AB = Generator(dim_A, dim_B).to(device)\ngen_BA = Generator(dim_B, dim_A).to(device)\ngen_opt = torch.optim.Adam(list(gen_AB.parameters()) + list(gen_BA.parameters()), lr=lr, betas=(0.5, 0.999))\ndisc_A = Discriminator(dim_A).to(device)\ndisc_A_opt = torch.optim.Adam(disc_A.parameters(), lr=lr, betas=(0.5, 0.999))\ndisc_B = Discriminator(dim_B).to(device)\ndisc_B_opt = torch.optim.Adam(disc_B.parameters(), lr=lr, betas=(0.5, 0.999))\n\ndef weights_init(m):\n    if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):\n        torch.nn.init.normal_(m.weight, 0.0, 0.02)\n    if isinstance(m, nn.BatchNorm2d):\n        torch.nn.init.normal_(m.weight, 0.0, 0.02)\n        torch.nn.init.constant_(m.bias, 0)","metadata":{"execution":{"iopub.status.busy":"2021-08-18T17:49:19.040799Z","iopub.status.idle":"2021-08-18T17:49:19.041386Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_disc_loss(real_X, fake_X, disc_X, adv_criterion):\n    disc_fake_X_hat = disc_X(fake_X.detach())\n    disc_fake_X_loss  = adv_criterion(disc_fake_X_hat,torch.zeros_like(disc_fake_X_hat))\n    disc_real_X_hat  = disc_X(real_X)\n    disc_real_X_loss  = adv_criterion(disc_real_X_hat,torch.ones_like(disc_real_X_hat))\n    disc_loss = (disc_real_X_loss +disc_fake_X_loss)/2\n    return disc_loss","metadata":{"execution":{"iopub.status.busy":"2021-08-18T17:49:19.0425Z","iopub.status.idle":"2021-08-18T17:49:19.043117Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_gen_adversarial_loss(real_X, disc_Y, gen_XY, adv_criterion):\n    fake_Y = gen_XY(real_X)\n    disc_fake_Y_hat  = disc_Y(fake_Y)\n    adversarial_loss  = adv_criterion(disc_fake_Y_hat,torch.ones_like(disc_fake_Y_hat))\n    return adversarial_loss, fake_Y","metadata":{"execution":{"iopub.status.busy":"2021-08-18T17:49:19.04434Z","iopub.status.idle":"2021-08-18T17:49:19.044951Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_identity_loss(real_X, gen_YX, identity_criterion):\n    identity_X = gen_YX(real_X)\n    identity_loss = identity_criterion(identity_X,real_X)\n    return identity_loss, identity_X","metadata":{"execution":{"iopub.status.busy":"2021-08-18T17:49:19.046125Z","iopub.status.idle":"2021-08-18T17:49:19.046729Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_cycle_consistency_loss(real_X, fake_Y, gen_YX, cycle_criterion):\n    cycle_X = gen_YX(fake_Y)\n    cycle_loss = cycle_criterion(cycle_X,real_X)\n    return cycle_loss, cycle_X","metadata":{"execution":{"iopub.status.busy":"2021-08-18T17:49:19.047769Z","iopub.status.idle":"2021-08-18T17:49:19.048351Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_gen_loss(real_A, real_B, gen_AB, gen_BA, disc_A, disc_B, adv_criterion, identity_criterion, cycle_criterion, lambda_identity=0.1, lambda_cycle=10):\n    adv_loss_BA, fake_A = get_gen_adversarial_loss(real_B, disc_A, gen_BA, adv_criterion)\n    adv_loss_AB, fake_B = get_gen_adversarial_loss(real_A, disc_B, gen_AB, adv_criterion)\n    gen_adversarial_loss = adv_loss_BA + adv_loss_AB\n    identity_loss_A, identity_A = get_identity_loss(real_A, gen_BA, identity_criterion)\n    identity_loss_B, identity_B = get_identity_loss(real_B, gen_AB, identity_criterion)\n    gen_identity_loss = identity_loss_A + identity_loss_B\n    cycle_loss_BA, cycle_A = get_cycle_consistency_loss(real_A, fake_B, gen_BA, cycle_criterion)\n    cycle_loss_AB, cycle_B = get_cycle_consistency_loss(real_B, fake_A, gen_AB, cycle_criterion)\n    gen_cycle_loss = cycle_loss_BA + cycle_loss_AB\n    gen_loss = lambda_identity * gen_identity_loss + lambda_cycle * gen_cycle_loss + gen_adversarial_loss\n    return gen_loss, fake_A, fake_B","metadata":{"execution":{"iopub.status.busy":"2021-08-18T17:49:19.049423Z","iopub.status.idle":"2021-08-18T17:49:19.050029Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def show_tensor_images(image_tensor, num_images=25, size=(1, 28, 28)):\n    image_tensor = (image_tensor + 1) / 2\n    image_shifted = image_tensor\n    image_unflat = image_shifted.detach().cpu().view(-1, *size)\n    image_grid = make_grid(image_unflat[:num_images], nrow=5)\n    plt.imshow(image_grid.permute(1, 2, 0).squeeze())\n    plt.show()\n\n","metadata":{"execution":{"iopub.status.busy":"2021-08-18T17:49:19.051155Z","iopub.status.idle":"2021-08-18T17:49:19.051777Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from skimage import color\nimport numpy as np\nplt.rcParams[\"figure.figsize\"] = (10, 10)\n\n\ndef train(save_model=False):\n    mean_generator_loss = 0\n    mean_discriminator_loss = 0\n    dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)\n    cur_step = 0\n\n    for epoch in range(n_epochs):\n        for real_A, real_B in tqdm(dataloader):\n            real_A = nn.functional.interpolate(real_A, size=target_shape)\n            real_B = nn.functional.interpolate(real_B, size=target_shape)\n            cur_batch_size = len(real_A)\n            real_A = real_A.to(device)\n            real_B = real_B.to(device)\n\n      \n            disc_A_opt.zero_grad() \n            with torch.no_grad():\n                fake_A = gen_BA(real_B)\n            disc_A_loss = get_disc_loss(real_A, fake_A, disc_A, adv_criterion)\n            disc_A_loss.backward(retain_graph=True)\n            disc_A_opt.step() \n\n           \n            disc_B_opt.zero_grad() \n            with torch.no_grad():\n                fake_B = gen_AB(real_A)\n            disc_B_loss = get_disc_loss(real_B, fake_B, disc_B, adv_criterion)\n            disc_B_loss.backward(retain_graph=True) \n            disc_B_opt.step()\n\n          \n            gen_opt.zero_grad()\n            gen_loss, fake_A, fake_B = get_gen_loss(\n                real_A, real_B, gen_AB, gen_BA, disc_A, disc_B, adv_criterion, recon_criterion, recon_criterion\n            )\n            gen_loss.backward() \n            gen_opt.step() \n            mean_discriminator_loss += disc_A_loss.item() / display_step\n            mean_generator_loss += gen_loss.item() / display_step\n            if cur_step % display_step == 0:\n                print(f\"Epoch {epoch}: Step {cur_step}: Generator (U-Net) loss: {mean_generator_loss}, Discriminator loss: {mean_discriminator_loss}\")\n                show_tensor_images(torch.cat([real_A, real_B]), size=(dim_A, target_shape, target_shape))\n                show_tensor_images(torch.cat([fake_B, fake_A]), size=(dim_B, target_shape, target_shape))\n                mean_generator_loss = 0\n                mean_discriminator_loss = 0\n                if save_model:\n                    torch.save({\n                        'gen_AB': gen_AB.state_dict(),\n                        'gen_BA': gen_BA.state_dict(),\n                        'gen_opt': gen_opt.state_dict(),\n                        'disc_A': disc_A.state_dict(),\n                        'disc_A_opt': disc_A_opt.state_dict(),\n                        'disc_B': disc_B.state_dict(),\n                        'disc_B_opt': disc_B_opt.state_dict()\n                    }, f\"cycleGAN_{cur_step}.pth\")\n            cur_step += 1\ntrain()","metadata":{"execution":{"iopub.status.busy":"2021-08-18T17:49:19.053002Z","iopub.status.idle":"2021-08-18T17:49:19.053627Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"torch.save(gen_AB.state_dict(),'./gen_AB.pt')\ntorch.save(gen_BA.state_dict(),'./gen_BA.pt')","metadata":{"execution":{"iopub.status.busy":"2021-08-18T17:49:19.054648Z","iopub.status.idle":"2021-08-18T17:49:19.055242Z"},"trusted":true},"execution_count":null,"outputs":[]}]}