{"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":"import torch\nimport torch.nn as nn \nimport numpy as np  \nfrom PIL import Image\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nimport os\nfrom torch.utils.data import Dataset, DataLoader  \nimport torch.optim as optim \nfrom tqdm import tqdm \nfrom torchvision.utils import save_image\nimport shutil","metadata":{"execution":{"iopub.status.busy":"2022-02-10T08:17:39.57058Z","iopub.execute_input":"2022-02-10T08:17:39.571099Z","iopub.status.idle":"2022-02-10T08:17:43.229828Z","shell.execute_reply.started":"2022-02-10T08:17:39.570893Z","shell.execute_reply":"2022-02-10T08:17:43.229103Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"torch.cuda.empty_cache()\ntorch.no_grad()","metadata":{"execution":{"iopub.status.busy":"2022-02-10T08:17:43.233048Z","iopub.execute_input":"2022-02-10T08:17:43.233501Z","iopub.status.idle":"2022-02-10T08:17:43.240609Z","shell.execute_reply.started":"2022-02-10T08:17:43.233474Z","shell.execute_reply":"2022-02-10T08:17:43.239976Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Configuration**","metadata":{}},{"cell_type":"code","source":"DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nTRAIN_DIR = \"../input/gan-getting-started/\"\nBATCH_SIZE = 1\nLEARNING_RATE = 1e-4 \nNUM_WORKERS = 0\nNUM_EPOCHS = 2\nLOAD_MODEL = False\nSAVE_MODEL = True\nCHECKPOINT_GEN_P = \"/kaggle/gen_P.pth.tar\"\nCHECKPOINT_GEN_M = \"/kaggle/gen_M.pth.tar\"\nCHECKPOINT_CRITIC_P = \"/kaggle/critic_P.pth.tar\"\nCHECKPOINT_CRITIC_M = \"/kaggle/critic_M.pth.tar\"\n\n#WGAN\nLAMBDA_GP = 100 \nLAMBDA_MEAN = 1 \nLAMBDA_cycle_MP = 1e-4 \nLAMBDA_CYCLE = 20\nLAMBDA_IDENTITY = 10 \nLAMBDA_W = 1e-4 \n\nprint(DEVICE)","metadata":{"execution":{"iopub.status.busy":"2022-02-10T08:17:43.242089Z","iopub.execute_input":"2022-02-10T08:17:43.24252Z","iopub.status.idle":"2022-02-10T08:17:43.299249Z","shell.execute_reply.started":"2022-02-10T08:17:43.242484Z","shell.execute_reply":"2022-02-10T08:17:43.298493Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"transforms = 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)\n\ntransforms_generate = A.Compose(\n    [\n        A.Resize(width=256, height=256),\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":"2022-02-10T08:17:43.301917Z","iopub.execute_input":"2022-02-10T08:17:43.302216Z","iopub.status.idle":"2022-02-10T08:17:43.310925Z","shell.execute_reply.started":"2022-02-10T08:17:43.302187Z","shell.execute_reply":"2022-02-10T08:17:43.309976Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Discriminator**","metadata":{}},{"cell_type":"code","source":"class Block(nn.Module):\n    def __init__(self, in_channels, out_channels, stride):\n        super().__init__()\n        self.conv = nn.Sequential(\n            nn.Conv2d(in_channels, out_channels, 4, stride, 1, bias = False, padding_mode = 'reflect'),\n            nn.InstanceNorm2d(out_channels), # try batchnorm\n            nn.LeakyReLU(0.2, inplace=True), # try relu\n        )\n    def forward(self, x):\n        return self.conv(x)","metadata":{"execution":{"iopub.status.busy":"2022-02-10T08:17:43.312381Z","iopub.execute_input":"2022-02-10T08:17:43.312811Z","iopub.status.idle":"2022-02-10T08:17:43.320423Z","shell.execute_reply.started":"2022-02-10T08:17:43.312775Z","shell.execute_reply":"2022-02-10T08:17:43.319593Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Discriminator(nn.Module):\n    def __init__(self, in_channels = 3, features = [64,128, 256, 512]):\n        super().__init__()\n        self.initial = nn.Sequential(\n            nn.Conv2d(in_channels, features[0], kernel_size = 4, stride = 2, padding = 1, padding_mode = 'reflect'),\n            nn.LeakyReLU(0.2, inplace=True), #try Relu\n        )\n        layers = []\n        in_channels = features[0]\n        \n        for feature in features[1:]:\n            layers.append(Block(in_channels, feature, stride = 1 if feature == features[-1] else 2))\n            in_channels = feature\n            \n        layers.append(nn.Conv2d(in_channels, 1, kernel_size=4, stride=1, padding=1, padding_mode='reflect')) # features[-1] = in_channels\n        self.model = nn.Sequential(*layers)\n        \n    def forward(self, x):\n        x = self.initial(x)\n        return self.model(x)\n        #return torch.sigmoid(self.model(x))","metadata":{"execution":{"iopub.status.busy":"2022-02-10T08:17:43.321714Z","iopub.execute_input":"2022-02-10T08:17:43.322173Z","iopub.status.idle":"2022-02-10T08:17:43.331475Z","shell.execute_reply.started":"2022-02-10T08:17:43.322138Z","shell.execute_reply":"2022-02-10T08:17:43.330656Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def test_discriminator():\n    x = torch.randn(1, 3, 256, 256)\n    model = Discriminator (in_channels=3)\n    preds = model(x) \n    print(preds.shape)\n\ntest_discriminator()","metadata":{"execution":{"iopub.status.busy":"2022-02-10T08:17:43.333781Z","iopub.execute_input":"2022-02-10T08:17:43.334216Z","iopub.status.idle":"2022-02-10T08:17:43.572572Z","shell.execute_reply.started":"2022-02-10T08:17:43.334179Z","shell.execute_reply":"2022-02-10T08:17:43.571869Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Generator**","metadata":{}},{"cell_type":"code","source":"class ConvBlock(nn.Module):\n    def __init__(self, in_channels, out_channels, down=True, use_act=True, **kwargs):\n        super().__init__()\n        self.conv = nn.Sequential(\n            nn.Conv2d(in_channels, out_channels, bias = False, padding_mode=\"reflect\", **kwargs) # bias = False\n            if down\n            else nn.ConvTranspose2d(in_channels, out_channels, bias = False, **kwargs), # bias = False\n            nn.InstanceNorm2d(out_channels),\n            nn.ReLU(inplace=True) if use_act else nn.Identity()\n        )\n\n    def forward(self, x):\n        return self.conv(x)\n\nclass ResidualBlock(nn.Module):\n    def __init__(self, channels):\n        super().__init__()\n        self.block = nn.Sequential(\n            ConvBlock(channels, channels, kernel_size=3, padding=1),\n            ConvBlock(channels, channels, use_act=False, kernel_size=3, padding=1),\n        )\n\n    def forward(self, x):\n        return x + self.block(x)","metadata":{"execution":{"iopub.status.busy":"2022-02-10T08:17:43.573826Z","iopub.execute_input":"2022-02-10T08:17:43.574214Z","iopub.status.idle":"2022-02-10T08:17:43.583101Z","shell.execute_reply.started":"2022-02-10T08:17:43.574178Z","shell.execute_reply":"2022-02-10T08:17:43.582377Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Generator(nn.Module):\n    def __init__(self, img_channels, num_features = 64, num_residuals=9):\n        super().__init__()\n        self.initial = nn.Sequential(\n            nn.Conv2d(img_channels, num_features, kernel_size=7, stride=1, padding=3, padding_mode=\"reflect\", bias = False), # bias = False\n            nn.InstanceNorm2d(num_features),\n            nn.ReLU(inplace=True),\n        )\n        self.down_blocks = nn.ModuleList(\n            [\n                ConvBlock(num_features, num_features*2, kernel_size=3, stride=2, padding=1),\n                ConvBlock(num_features*2, num_features*4, kernel_size=3, stride=2, padding=1),\n            ]\n        )\n        self.res_blocks = nn.Sequential(\n            *[ResidualBlock(num_features*4) for _ in range(num_residuals)]\n        )\n        self.up_blocks = nn.ModuleList(\n            [\n                ConvBlock(num_features*4, num_features*2, down=False, kernel_size=3, stride=2, padding=1, output_padding=1),\n                ConvBlock(num_features*2, num_features*1, down=False, kernel_size=3, stride=2, padding=1, output_padding=1),\n            ]\n        )\n\n        self.last = nn.Conv2d(num_features*1, img_channels, kernel_size=7, stride=1, padding=3, padding_mode=\"reflect\")\n\n    def forward(self, x):\n        x = self.initial(x)\n        for layer in self.down_blocks:\n            x = layer(x)\n        x = self.res_blocks(x)\n        for layer in self.up_blocks:\n            x = layer(x)\n        return torch.tanh(self.last(x))","metadata":{"execution":{"iopub.status.busy":"2022-02-10T08:17:43.584296Z","iopub.execute_input":"2022-02-10T08:17:43.584549Z","iopub.status.idle":"2022-02-10T08:17:43.597539Z","shell.execute_reply.started":"2022-02-10T08:17:43.584513Z","shell.execute_reply":"2022-02-10T08:17:43.596898Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def test_generator():\n    img_channels = 3\n    img_size = 256\n    x = torch.randn((2, img_channels, img_size, img_size))\n    gen = Generator(img_channels, 9)\n    print(gen(x).shape)\n\ntest_generator()","metadata":{"execution":{"iopub.status.busy":"2022-02-10T08:17:43.600592Z","iopub.execute_input":"2022-02-10T08:17:43.601031Z","iopub.status.idle":"2022-02-10T08:17:43.883597Z","shell.execute_reply.started":"2022-02-10T08:17:43.601Z","shell.execute_reply":"2022-02-10T08:17:43.882738Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Saving and loading model**","metadata":{}},{"cell_type":"code","source":"def save_checkpoint(model, optimizer, filename = './mycheckpoint.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)","metadata":{"execution":{"iopub.status.busy":"2022-02-10T08:17:43.884879Z","iopub.execute_input":"2022-02-10T08:17:43.885137Z","iopub.status.idle":"2022-02-10T08:17:43.890652Z","shell.execute_reply.started":"2022-02-10T08:17:43.8851Z","shell.execute_reply":"2022-02-10T08:17:43.889702Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def load_checkpoint(checkpoint_file, model, optimizer, lr):\n    print('===> Loading checkpoint ....')\n    checkpoint = torch.load(checkpoint_file, map_location = DEVICE)\n    model.load_state_dict(checkpoint['state_dict'])\n    optimizer.load_state_dict(checkpoint['optimizer'])\n    # override lr of previous checkpoint\n    for param_group in optimizer.param_groups:\n        param_group['lr'] = lr","metadata":{"execution":{"iopub.status.busy":"2022-02-10T08:17:43.891786Z","iopub.execute_input":"2022-02-10T08:17:43.892401Z","iopub.status.idle":"2022-02-10T08:17:43.899022Z","shell.execute_reply.started":"2022-02-10T08:17:43.892363Z","shell.execute_reply":"2022-02-10T08:17:43.898227Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Dataset** ","metadata":{}},{"cell_type":"code","source":"class MonetPictureDataset(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        self.length_dataset = max(len(self.monet_images), len(self.picture_images)) # 1000, 1500\n        self.monet_len = len(self.monet_images)\n        self.picture_len = len(self.picture_images)\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            augmentations = self.transform(image=monet_img, image0=picture_img)\n            monet_img = augmentations[\"image\"]\n            picture_img = augmentations[\"image0\"]\n\n        return monet_img, picture_img","metadata":{"execution":{"iopub.status.busy":"2022-02-10T08:17:43.90015Z","iopub.execute_input":"2022-02-10T08:17:43.900409Z","iopub.status.idle":"2022-02-10T08:17:43.91158Z","shell.execute_reply.started":"2022-02-10T08:17:43.900372Z","shell.execute_reply":"2022-02-10T08:17:43.910568Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Training**","metadata":{}},{"cell_type":"code","source":"def gradient_penalty(critic, real, fake, device=\"cpu\"):\n    with torch.cuda.amp.autocast(): # necessqry for float 16 (more fast)\n        BATCH_SIZE, C, H, W = real.shape\n        alpha = torch.rand((BATCH_SIZE, 1, 1, 1)).repeat(1, C, H, W).to(device)\n        interpolated_images = real * alpha + fake * (1 - alpha)\n\n        # Calculate critic scores\n        mixed_scores = critic(interpolated_images)\n\n        # Take the gradient of the scores with respect to the images\n        gradient = torch.autograd.grad(\n          inputs=interpolated_images,\n          outputs=mixed_scores,\n          grad_outputs=torch.ones_like(mixed_scores),\n          create_graph=False,\n          retain_graph=False,\n        )[0]\n        gradient = gradient.view(gradient.shape[0], -1)\n        gradient_norm = gradient.norm(2, dim=1)\n        gradient_penalty = torch.mean((gradient_norm - 1) ** 2)\n        return gradient_penalty","metadata":{"execution":{"iopub.status.busy":"2022-02-10T08:17:43.91295Z","iopub.execute_input":"2022-02-10T08:17:43.913465Z","iopub.status.idle":"2022-02-10T08:17:43.921194Z","shell.execute_reply.started":"2022-02-10T08:17:43.913431Z","shell.execute_reply":"2022-02-10T08:17:43.920476Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def train_fn(disc_P, disc_M, gen_M, gen_P, loader, opt_disc, opt_gen, l1, mse, d_scaler, g_scaler):\n    P_reals = 0\n    P_fakes = 0\n    Gen_Loss = 0\n    Disc_Loss = 0 \n    loop = tqdm(loader, leave=True)\n\n    for idx, (monet, picture) in enumerate(loop):\n        monet = monet.to(DEVICE)\n        picture = picture.to(DEVICE)\n\n        # Train Discriminators H and Z\n        with torch.cuda.amp.autocast(): # necessqry for float 16 (more fast)\n            fake_picture = gen_P(monet)\n            D_P_real = disc_P(picture)\n            D_P_fake = disc_P(fake_picture.detach()) # we gonna use fake_picture later on when we train generator\n\n            P_reals += D_P_real.mean().item()\n            P_fakes += D_P_fake.mean().item()\n\n            D_P_real_loss = mse(D_P_real, torch.ones_like(D_P_real))\n            D_P_fake_loss = mse(D_P_fake, torch.zeros_like(D_P_fake))\n\n            gp_P = gradient_penalty(disc_P,picture,fake_picture, device = DEVICE)\n            D_P_loss_W = -(torch.mean(D_P_real.view(-1))-torch.mean(D_P_fake.view(-1))) + LAMBDA_GP*gp_P #reshape\n\n            D_P_loss = D_P_real_loss + D_P_fake_loss + LAMBDA_W * D_P_loss_W\n\n            fake_monet = gen_M(picture)\n            D_M_real = disc_M(monet)\n            D_M_fake = disc_M(fake_monet.detach())\n            D_M_real_loss = mse(D_M_real, torch.ones_like(D_M_real))\n            D_M_fake_loss = mse(D_M_fake, torch.zeros_like(D_M_fake))\n\n            gp_M = gradient_penalty(disc_M,monet,fake_monet, device = DEVICE)\n            D_M_loss_W = -(torch.mean(D_M_real.view(-1))-torch.mean(D_M_fake.view(-1))) + LAMBDA_GP*gp_M \n            \n            D_M_loss = D_M_real_loss + D_M_fake_loss + LAMBDA_W * D_M_loss_W\n\n            # put it togethor\n            D_loss = (D_P_loss + D_M_loss)/2 # try modify it\n             \n            Disc_Loss += D_loss.item()\n\n        opt_disc.zero_grad()\n        d_scaler.scale(D_loss).backward(retain_graph=True) #\n        d_scaler.step(opt_disc)\n        d_scaler.update()\n\n        # Train Generators H and Z\n        with torch.cuda.amp.autocast():\n            # adversarial loss for both generators\n            D_P_fake = disc_P(fake_picture)\n            D_M_fake = disc_M(fake_monet)\n            \n            loss_G_P_W = -torch.mean(D_P_fake)\n            loss_G_M_W = -torch.mean(D_M_fake)\n\n            loss_G_P = mse(D_P_fake, torch.ones_like(D_P_fake)) + LAMBDA_MEAN* loss_G_P_W\n            loss_G_M = mse(D_M_fake, torch.ones_like(D_M_fake)) + LAMBDA_MEAN* loss_G_M_W\n            \n            # cycle loss\n            cycle_monet = gen_M(fake_picture)\n            cycle_picture = gen_P(fake_monet)\n\n            gp_gen_P = gradient_penalty(gen_P,picture,cycle_picture, device = DEVICE)\n            gp_gen_M = gradient_penalty(gen_M,monet,cycle_monet, device = DEVICE)\n\n            cycle_monet_loss = l1(monet, cycle_monet) + LAMBDA_cycle_MP * gp_gen_M\n            cycle_picture_loss = l1(picture, cycle_picture) + LAMBDA_cycle_MP * gp_gen_P\n\n            # identity loss (remove these for efficiency if you set lambda_identity=0)\n            identity_monet = gen_M(monet)\n            identity_picture = gen_P(picture)\n            identity_monet_loss = l1(monet, identity_monet)\n            identity_picture_loss = l1(picture, identity_picture)\n\n            # add all togethor\n            G_loss = (\n                loss_G_M\n                + loss_G_P\n                + cycle_monet_loss * LAMBDA_CYCLE\n                + cycle_picture_loss * LAMBDA_CYCLE\n                + identity_picture_loss * LAMBDA_IDENTITY\n                + identity_monet_loss * LAMBDA_IDENTITY\n            )\n            \n            Gen_Loss += G_loss.item()\n\n        opt_gen.zero_grad()\n        g_scaler.scale(G_loss).backward(retain_graph = True)\n        g_scaler.step(opt_gen)\n        g_scaler.update() \n        \n        loop.set_postfix(P_real=P_reals/(idx+1), P_fake=P_fakes/(idx+1), G_L = Gen_Loss/(idx+1), D_L = Disc_Loss/(idx+1) )","metadata":{"execution":{"iopub.status.busy":"2022-02-10T08:17:43.922567Z","iopub.execute_input":"2022-02-10T08:17:43.923051Z","iopub.status.idle":"2022-02-10T08:17:43.943605Z","shell.execute_reply.started":"2022-02-10T08:17:43.923016Z","shell.execute_reply":"2022-02-10T08:17:43.942899Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def main(disc_P, disc_M, gen_P, gen_M, opt_disc, opt_gen):\n\n    #WGAN end\n\n    L1 = nn.L1Loss()\n    mse = nn.MSELoss()\n\n    if LOAD_MODEL:\n        load_checkpoint(\n            CHECKPOINT_GEN_P, gen_P, opt_gen, LEARNING_RATE,\n        )\n        load_checkpoint(\n            CHECKPOINT_GEN_M, gen_M, opt_gen, LEARNING_RATE,\n        )\n        load_checkpoint(\n            CHECKPOINT_CRITIC_P, disc_P, opt_disc, LEARNING_RATE,\n        )\n        load_checkpoint(\n            CHECKPOINT_CRITIC_M, disc_M, opt_disc, LEARNING_RATE,\n        )\n\n    dataset = MonetPictureDataset(\n        root_picture=TRAIN_DIR+\"/photo_jpg\", root_monet=TRAIN_DIR+\"/monet_jpg\", transform=transforms\n    )\n    loader = DataLoader(\n        dataset,\n        batch_size=BATCH_SIZE,\n        shuffle=True,\n        num_workers=NUM_WORKERS,\n        pin_memory=True\n    )\n    g_scaler = torch.cuda.amp.GradScaler()\n    d_scaler = torch.cuda.amp.GradScaler()\n\n    for epoch in range(NUM_EPOCHS):\n        train_fn(disc_P, disc_M, gen_M, gen_P, loader, opt_disc, opt_gen, L1, mse, d_scaler, g_scaler)\n\n        if SAVE_MODEL:\n            save_checkpoint(gen_P, opt_gen, filename=CHECKPOINT_GEN_P)\n            save_checkpoint(gen_M, opt_gen, filename=CHECKPOINT_GEN_M)\n            save_checkpoint(disc_P, opt_disc, filename=CHECKPOINT_CRITIC_P)\n            save_checkpoint(disc_M, opt_disc, filename=CHECKPOINT_CRITIC_M)","metadata":{"execution":{"iopub.status.busy":"2022-02-10T08:17:43.944771Z","iopub.execute_input":"2022-02-10T08:17:43.945215Z","iopub.status.idle":"2022-02-10T08:17:43.956591Z","shell.execute_reply.started":"2022-02-10T08:17:43.94518Z","shell.execute_reply":"2022-02-10T08:17:43.955898Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"disc_P = Discriminator(in_channels=3).to(DEVICE)\ndisc_M = Discriminator(in_channels=3).to(DEVICE)\ngen_M = Generator(img_channels=3, num_residuals=9).to(DEVICE)\ngen_P = Generator(img_channels=3, num_residuals=9).to(DEVICE)\n\nopt_disc = optim.Adam(\nlist(disc_P.parameters()) + list(disc_M.parameters()),\nlr= LEARNING_RATE,\nbetas=(0.5, 0.999),\n)\n\nopt_gen = optim.Adam(\nlist(gen_M.parameters()) + list(gen_P.parameters()),\nlr=LEARNING_RATE,\nbetas=(0.5, 0.999),\n)","metadata":{"execution":{"iopub.status.busy":"2022-02-10T08:17:43.958892Z","iopub.execute_input":"2022-02-10T08:17:43.959476Z","iopub.status.idle":"2022-02-10T08:17:47.229846Z","shell.execute_reply.started":"2022-02-10T08:17:43.959427Z","shell.execute_reply":"2022-02-10T08:17:47.229116Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"main(disc_P, disc_M, gen_P, gen_M, opt_disc, opt_gen)","metadata":{"execution":{"iopub.status.busy":"2022-02-10T08:17:47.231295Z","iopub.execute_input":"2022-02-10T08:17:47.231547Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Load trained model**","metadata":{}},{"cell_type":"code","source":"load_checkpoint(\nCHECKPOINT_GEN_P, gen_P, opt_gen, LEARNING_RATE,\n)\nload_checkpoint(\nCHECKPOINT_GEN_M, gen_M, opt_gen, LEARNING_RATE,\n)\nload_checkpoint(\nCHECKPOINT_CRITIC_P, disc_P, opt_disc, LEARNING_RATE,\n)\nload_checkpoint(\nCHECKPOINT_CRITIC_M, disc_M, opt_disc, LEARNING_RATE,\n)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Generate images**","metadata":{}},{"cell_type":"code","source":"dataset = MonetPictureDataset(\n        root_picture=TRAIN_DIR+\"/photo_jpg\", root_monet=TRAIN_DIR+\"/monet_jpg\", transform=transforms_generate\n    ) \nloader = DataLoader(\n    dataset,\n    batch_size=BATCH_SIZE,\n    shuffle=True,\n    num_workers=NUM_WORKERS,\n    pin_memory=True\n)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"disc_P.eval()\ndisc_M.eval()\ngen_M.eval()\ngen_P.eval()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if not os.path.exists('/kaggle/images'):\n    os.makedirs('/kaggle/images')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for idx, (monet, picture) in enumerate(loader):\n    with torch.no_grad():\n        fake_picture = gen_M(picture.to(DEVICE))\n        save_image(fake_picture*0.5+0.5, f'/kaggle/images/picture_{idx}.jpg')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"shutil.make_archive(\"/kaggle/working/images\", 'zip', \"/kaggle/images\")\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}