{"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":"# Imports","metadata":{"execution":{"iopub.status.busy":"2022-09-22T03:29:04.782528Z","iopub.execute_input":"2022-09-22T03:29:04.783201Z","iopub.status.idle":"2022-09-22T03:29:04.787474Z","shell.execute_reply.started":"2022-09-22T03:29:04.783163Z","shell.execute_reply":"2022-09-22T03:29:04.786292Z"}}},{"cell_type":"code","source":"import numpy as np\nfrom PIL import Image\nfrom tqdm import tqdm\nfrom pathlib import Path\nimport matplotlib.pyplot as plt\n\nimport torch\nfrom torch.utils.data import Dataset, DataLoader, random_split\n\nfrom torchvision import transforms\nfrom torchvision.utils import make_grid","metadata":{"execution":{"iopub.status.busy":"2022-09-22T15:10:33.951837Z","iopub.execute_input":"2022-09-22T15:10:33.952204Z","iopub.status.idle":"2022-09-22T15:10:33.959161Z","shell.execute_reply.started":"2022-09-22T15:10:33.952172Z","shell.execute_reply":"2022-09-22T15:10:33.957202Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Paths\nDefining the paths used for the dataset","metadata":{}},{"cell_type":"code","source":"monet_path = Path(\"../input/gan-getting-started/monet_jpg\")\nphoto_path = Path(\"../input/gan-getting-started/photo_jpg\")","metadata":{"execution":{"iopub.status.busy":"2022-09-22T15:10:34.597428Z","iopub.execute_input":"2022-09-22T15:10:34.597807Z","iopub.status.idle":"2022-09-22T15:10:34.603357Z","shell.execute_reply.started":"2022-09-22T15:10:34.597774Z","shell.execute_reply":"2022-09-22T15:10:34.602318Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Checking the number of images in each subfolder","metadata":{}},{"cell_type":"code","source":"len(list(monet_path.glob(\"**/*.jpg\"))), len(list(photo_path.glob(\"**/*.jpg\")))","metadata":{"execution":{"iopub.status.busy":"2022-09-22T15:10:35.054827Z","iopub.execute_input":"2022-09-22T15:10:35.055458Z","iopub.status.idle":"2022-09-22T15:10:37.384696Z","shell.execute_reply.started":"2022-09-22T15:10:35.055422Z","shell.execute_reply":"2022-09-22T15:10:37.383670Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Checking an image simple from each class and to give a feel of the generated image we show a blended version of the two images. We calculate the blend as follow:\n$$ B = (1-\\alpha) \\times X + \\alpha \\times Y$$","metadata":{}},{"cell_type":"code","source":"sample_monet, sample_photo = Image.open(list(monet_path.glob(\"**/*.jpg\"))[0]), Image.open(list(photo_path.glob(\"**/*.jpg\"))[0])","metadata":{"execution":{"iopub.status.busy":"2022-09-22T15:10:37.386795Z","iopub.execute_input":"2022-09-22T15:10:37.387467Z","iopub.status.idle":"2022-09-22T15:10:37.467107Z","shell.execute_reply.started":"2022-09-22T15:10:37.387429Z","shell.execute_reply":"2022-09-22T15:10:37.466276Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"fig, ax = plt.subplots(1, 3, figsize=(20, 4))\nax[0].imshow(sample_monet)\nax[0].set_title(\"Monet\")\nax[0].set(xticks=[], yticks=[])\nax[1].imshow(sample_photo)\nax[1].set_title(\"Photo\")\nax[1].set(xticks=[], yticks=[])\nax[2].imshow(Image.blend(sample_monet, sample_photo, 0.6))\nax[2].set_title(\"Blended\")\nax[2].set(xticks=[], yticks=[])\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2022-09-22T15:10:37.468201Z","iopub.execute_input":"2022-09-22T15:10:37.468861Z","iopub.status.idle":"2022-09-22T15:10:37.776867Z","shell.execute_reply.started":"2022-09-22T15:10:37.468833Z","shell.execute_reply":"2022-09-22T15:10:37.775713Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Check the image size and format","metadata":{}},{"cell_type":"code","source":"Image.open(list(photo_path.glob(\"**/*.jpg\"))[0]).size, Image.open(list(photo_path.glob(\"**/*.jpg\"))[0]).mode","metadata":{"execution":{"iopub.status.busy":"2022-09-22T15:10:37.779542Z","iopub.execute_input":"2022-09-22T15:10:37.780630Z","iopub.status.idle":"2022-09-22T15:10:37.934240Z","shell.execute_reply.started":"2022-09-22T15:10:37.780560Z","shell.execute_reply":"2022-09-22T15:10:37.933124Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"Image.open(list(monet_path.glob(\"**/*.jpg\"))[0]).size, Image.open(list(monet_path.glob(\"**/*.jpg\"))[0]).mode","metadata":{"execution":{"iopub.status.busy":"2022-09-22T15:10:37.935546Z","iopub.execute_input":"2022-09-22T15:10:37.936089Z","iopub.status.idle":"2022-09-22T15:10:37.952249Z","shell.execute_reply.started":"2022-09-22T15:10:37.936048Z","shell.execute_reply":"2022-09-22T15:10:37.951418Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Dataset\nCreating a PyTorch custom dataset","metadata":{}},{"cell_type":"code","source":"class MonetPhoto(Dataset):\n    def __init__(self, monet_path: Path = Path(\"../input/gan-getting-started/monet_jpg\"), photo_path:Path = Path(\"../input/gan-getting-started/photo_jpg\"), transform=None):\n        self.monet_path = monet_path\n        self.monet_images = list(monet_path.glob(\"**/*.jpg\"))\n        self.photo_path = photo_path\n        self.photo_images = list(photo_path.glob(\"**/*.jpg\"))\n        self.transform = transform\n\n    def __len__(self):\n        return max(len(self.monet_images), len(self.photo_images))\n\n    def __getitem__(self, idx):\n        monet_image = Image.open(self.monet_images[idx % len(self.monet_images)])\n        photo_image = Image.open(self.photo_images[idx % len(self.photo_images)])\n        if self.transform:\n            monet_image = self.transform(monet_image)\n            photo_image = self.transform(photo_image)\n        return monet_image, photo_image","metadata":{"execution":{"iopub.status.busy":"2022-09-22T15:10:37.953490Z","iopub.execute_input":"2022-09-22T15:10:37.954202Z","iopub.status.idle":"2022-09-22T15:10:37.963695Z","shell.execute_reply.started":"2022-09-22T15:10:37.954165Z","shell.execute_reply":"2022-09-22T15:10:37.962705Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Checking the dataset.\n1. Note that we use PIL to open the pictures so we need to transforme them into a tensor\n2. In addition an image size of (256, 256) migth be too big for a baseline so we resize it to (32, 32)","metadata":{}},{"cell_type":"code","source":"dataset = MonetPhoto(transform=transforms.Compose([transforms.ToTensor(), transforms.Normalize(0.5, 0.5)]))\n# dataset = MonetPhoto(transform=transforms.Compose([transforms.Resize((32, 32)), transforms.ToTensor(), transforms.Normalize(0.5, 0.5)]))","metadata":{"execution":{"iopub.status.busy":"2022-09-22T15:10:37.964949Z","iopub.execute_input":"2022-09-22T15:10:37.965201Z","iopub.status.idle":"2022-09-22T15:10:38.045376Z","shell.execute_reply.started":"2022-09-22T15:10:37.965177Z","shell.execute_reply":"2022-09-22T15:10:38.044513Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"We check the shape of one simple","metadata":{}},{"cell_type":"code","source":"x, y = dataset.__getitem__(1)\nx.shape, y.shape","metadata":{"execution":{"iopub.status.busy":"2022-09-22T15:10:38.046601Z","iopub.execute_input":"2022-09-22T15:10:38.047316Z","iopub.status.idle":"2022-09-22T15:10:38.071572Z","shell.execute_reply.started":"2022-09-22T15:10:38.047277Z","shell.execute_reply":"2022-09-22T15:10:38.070612Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Note that the images has been normalized into a range of [-1, 1] since we will use a tanh activation on the output of the generator","metadata":{}},{"cell_type":"code","source":"x.max(), x.min(), y.max(), y.min()","metadata":{"execution":{"iopub.status.busy":"2022-09-22T15:10:38.072887Z","iopub.execute_input":"2022-09-22T15:10:38.073197Z","iopub.status.idle":"2022-09-22T15:10:38.083293Z","shell.execute_reply.started":"2022-09-22T15:10:38.073164Z","shell.execute_reply":"2022-09-22T15:10:38.082180Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"We split the dataset into a train and test dataset with a ratio of 80%-20%","metadata":{}},{"cell_type":"code","source":"train_size, test_size = int(len(dataset) * 0.8), len(dataset) - int(len(dataset) * 0.8)","metadata":{"execution":{"iopub.status.busy":"2022-09-22T15:10:38.086850Z","iopub.execute_input":"2022-09-22T15:10:38.087195Z","iopub.status.idle":"2022-09-22T15:10:38.091511Z","shell.execute_reply.started":"2022-09-22T15:10:38.087158Z","shell.execute_reply":"2022-09-22T15:10:38.090182Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"training_data, testing_data = random_split(dataset, [train_size, test_size])","metadata":{"execution":{"iopub.status.busy":"2022-09-22T15:10:38.092894Z","iopub.execute_input":"2022-09-22T15:10:38.095014Z","iopub.status.idle":"2022-09-22T15:10:38.102758Z","shell.execute_reply.started":"2022-09-22T15:10:38.094971Z","shell.execute_reply":"2022-09-22T15:10:38.101701Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"We create the dataloaders\n* This is not a classification problem so shuffling the data has no real meaning here, it's up to you","metadata":{}},{"cell_type":"code","source":"train_dataloader = DataLoader(training_data, batch_size=32, shuffle=True, num_workers=2)\ntest_dataloader = DataLoader(testing_data, batch_size=32, shuffle=True, num_workers=2)","metadata":{"execution":{"iopub.status.busy":"2022-09-22T15:10:38.270234Z","iopub.execute_input":"2022-09-22T15:10:38.270513Z","iopub.status.idle":"2022-09-22T15:10:38.275177Z","shell.execute_reply.started":"2022-09-22T15:10:38.270487Z","shell.execute_reply":"2022-09-22T15:10:38.274103Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"next(iter(test_dataloader))[0][:32].shape","metadata":{"execution":{"iopub.status.busy":"2022-09-22T15:10:38.429976Z","iopub.execute_input":"2022-09-22T15:10:38.430257Z","iopub.status.idle":"2022-09-22T15:10:39.535427Z","shell.execute_reply.started":"2022-09-22T15:10:38.430231Z","shell.execute_reply":"2022-09-22T15:10:39.534389Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for x, y in train_dataloader:\n    print(x.size(), y.size())\n    break","metadata":{"execution":{"iopub.status.busy":"2022-09-22T15:10:39.537939Z","iopub.execute_input":"2022-09-22T15:10:39.538844Z","iopub.status.idle":"2022-09-22T15:10:40.486935Z","shell.execute_reply.started":"2022-09-22T15:10:39.538800Z","shell.execute_reply":"2022-09-22T15:10:40.485720Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model\nFor our baseline we are going to use a simple CNN model for the generator and the discriminator. The training process will follow the original GAN process without the cycle process. In the future we migth use a diffrent approch.","metadata":{}},{"cell_type":"code","source":"device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\ndevice","metadata":{"execution":{"iopub.status.busy":"2022-09-22T15:10:40.488910Z","iopub.execute_input":"2022-09-22T15:10:40.489283Z","iopub.status.idle":"2022-09-22T15:10:40.553406Z","shell.execute_reply.started":"2022-09-22T15:10:40.489245Z","shell.execute_reply":"2022-09-22T15:10:40.552377Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Generator","metadata":{}},{"cell_type":"code","source":"class Generator(torch.nn.Module):\n    def __init__(\n        self,\n        in_channels: int = 50,\n        out_channels: int = 3,\n        nb_features: int = 32,\n        nb_layers: int = 2,\n        device=\"cuda\",\n    ):\n        super(Generator, self).__init__()\n        self.device = device\n        self.blocks = torch.nn.ModuleList(\n            [\n                self._block(\n                    in_channels=in_channels,\n                    out_channels=nb_features * (2**0),\n                    kernel_size=(4, 4),\n                    stride=2,\n                    padding=1,\n                    bias=False,\n                    down=True\n                )\n            ]\n        )\n        for i in range(nb_layers):\n            self.blocks.append(\n                self._block(\n                    in_channels=nb_features * (2**i),\n                    out_channels=nb_features * (2 ** (i + 1)),\n                    kernel_size=(4, 4),\n                    stride=2,\n                    padding=1,\n                    down=True\n                )\n            )\n        for i in range(nb_layers, 0, -1):\n            self.blocks.append(\n                self._block(\n                    in_channels=nb_features * (2 ** (i)),\n                    out_channels=nb_features * (2 ** (i - 1)),\n                    kernel_size=(2, 2),\n                    stride=2,\n                    bias=False,\n                )\n            )\n        self.blocks.append(\n            self._block(\n                in_channels=nb_features,\n                out_channels=out_channels,\n                kernel_size=(2, 2),\n                stride=2,\n                activation=False,\n                bias=True,\n            )\n        )\n\n    def _block(\n        self,\n        in_channels: int,\n        out_channels: int,\n        kernel_size: tuple = (4, 4),\n        stride: int = 1,\n        padding: int = 0,\n        padding_mode: str = \"reflect\",\n        bias: bool = False,\n        activation: bool = True,\n        down: bool = False\n    ):\n        return torch.nn.Sequential(\n            torch.nn.Conv2d(\n                in_channels=in_channels,\n                out_channels=out_channels,\n                kernel_size=kernel_size,\n                stride=stride,\n                padding=padding,\n                padding_mode=padding_mode,\n                bias=bias,\n                device=self.device,\n            ) if down else torch.nn.ConvTranspose2d(\n                in_channels=in_channels,\n                out_channels=out_channels,\n                kernel_size=kernel_size,\n                stride=stride,\n                bias=bias,\n                device=self.device,\n            ),\n            torch.nn.Identity(device=self.device)\n            if bias\n            else torch.nn.BatchNorm2d(\n                num_features=out_channels, eps=0.8, device=self.device\n            ),\n            torch.nn.LeakyReLU(0.2) if activation else torch.nn.Tanh(),\n        )\n\n    def forward(self, x: torch.Tensor) -> torch.Tensor:\n        for block in self.blocks:\n            x = block(x)\n        return x","metadata":{"execution":{"iopub.status.busy":"2022-09-22T15:10:40.556522Z","iopub.execute_input":"2022-09-22T15:10:40.557177Z","iopub.status.idle":"2022-09-22T15:10:40.572184Z","shell.execute_reply.started":"2022-09-22T15:10:40.557138Z","shell.execute_reply":"2022-09-22T15:10:40.571148Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Discriminator","metadata":{}},{"cell_type":"code","source":"class Discriminator(torch.nn.Module):\n    def __init__(\n        self,\n        in_channels: int = 3,\n        out_channels: int = 1,\n        nb_features: int = 32,\n        nb_layers: int = 2,\n        device=\"cuda\",\n    ):\n        super(Discriminator, self).__init__()\n        self.device = device\n        self.blocks = torch.nn.ModuleList(\n            [\n                self._block(\n                    in_channels=in_channels,\n                    out_channels=nb_features,\n                    kernel_size=(4, 4),\n                    stride=2,\n                    normalize=False,\n                )\n            ]\n        )\n        for i in range(nb_layers):\n            self.blocks.append(\n                self._block(\n                    in_channels=nb_features * (2**i),\n                    out_channels=nb_features * (2 ** (i + 1)),\n                    kernel_size=(4, 4),\n                    stride=2,\n                    padding=1,\n                )\n            )\n        self.blocks.append(\n            self._block(\n                in_channels=nb_features * (2**nb_layers),\n                out_channels=out_channels,\n                kernel_size=(2, 2),\n                stride=2,\n                padding=0,\n                normalize=False,\n            )\n        )\n\n    def _block(\n        self,\n        in_channels: int,\n        out_channels: int,\n        kernel_size: tuple = (4, 4),\n        stride: int = 1,\n        padding: int = 0,\n        padding_mode: str = \"reflect\",\n        bias: bool = False,\n        normalize: bool = True,\n    ):\n        return torch.nn.Sequential(\n            torch.nn.Conv2d(\n                in_channels=in_channels,\n                out_channels=out_channels,\n                kernel_size=kernel_size,\n                stride=stride,\n                padding=padding,\n                padding_mode=padding_mode,\n                bias=bias,\n                device=self.device,\n            ),\n            torch.nn.BatchNorm2d(num_features=out_channels, eps=0.8, device=self.device)\n            if not bias and normalize\n            else torch.nn.Identity(device=self.device),\n            torch.nn.LeakyReLU(0.2),\n        )\n\n    def forward(self, x: torch.Tensor) -> torch.Tensor:\n        for block in self.blocks:\n            x = block(x)\n        return x.view(x.size(0), -1)","metadata":{"execution":{"iopub.status.busy":"2022-09-22T15:10:40.573961Z","iopub.execute_input":"2022-09-22T15:10:40.574614Z","iopub.status.idle":"2022-09-22T15:10:40.589107Z","shell.execute_reply.started":"2022-09-22T15:10:40.574569Z","shell.execute_reply":"2022-09-22T15:10:40.588180Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Sanitary check","metadata":{}},{"cell_type":"code","source":"generator = Generator(in_channels=3, nb_layers=4, device=device)\ndiscriminator = Discriminator(nb_layers=5, device=device)\nwith torch.no_grad():\n    fake_img = generator(torch.randn((1, 3, 256, 256), device=device))\n    print(fake_img.shape)\n    print(discriminator(fake_img).shape)","metadata":{"execution":{"iopub.status.busy":"2022-09-22T15:10:40.591689Z","iopub.execute_input":"2022-09-22T15:10:40.592304Z","iopub.status.idle":"2022-09-22T15:10:48.756384Z","shell.execute_reply.started":"2022-09-22T15:10:40.592266Z","shell.execute_reply":"2022-09-22T15:10:48.754603Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Trainer\nI usually like to create a class handling the training, logging and model loading but this is just a preference.","metadata":{}},{"cell_type":"code","source":"class Trainer:\n    def __init__(\n        self,\n        train_data: DataLoader,\n        val_data: DataLoader,\n        generator: torch.nn.Module,\n        discriminator: torch.nn.Module,\n        nb_epochs: int = 5,\n        device: str = \"cuda\",\n        save_path: str = None,\n    ):\n        self.train_data = train_data\n        self.val_data = val_data\n        self.generator = generator\n        self.discriminator = discriminator\n        self.nb_epochs = nb_epochs\n        self.device = device\n        self.save_path = Path(save_path) if save_path else save_path\n\n        self.z = next(iter(self.val_data))[0][:32].to(self.device)\n        \n        self.logs = {\n            \"Step\": [],\n            \"Train_g_loss\": [],\n            \"Train_d_loss\": [],\n            \"Val_g_loss\": [],\n            \"Val_d_loss\": [],\n            \"Samples\": [],\n        }\n\n    def init_optimizers(self, lr: float=1e-3, betas: tuple=(0.5, 0.999)):\n        self.g_optimizer = torch.optim.Adam(\n            self.generator.parameters(), lr=lr, betas=betas\n        )\n        self.d_optimizer = torch.optim.Adam(\n            self.discriminator.parameters(), lr=lr, betas=betas\n        )\n\n    def train(self):\n        assert (self.g_optimizer is not None) and (\n            self.d_optimizer is not None\n        ), \"Please run Trainer().init_optimizer()\"\n\n        # Loading best model\n        if self.save_path and self.save_path.exists():\n            self.load_model()\n            \n        best_score = torch.inf\n\n        for i in range(self.nb_epochs):\n            train_d_loss, train_g_loss, val_d_loss, val_g_loss = 0, 0, 0, 0\n            self.generator.train()\n            self.discriminator.train()\n            # Train loop\n            loop = tqdm(\n                enumerate(self.train_data),\n                desc=f\"Epoch {i + 1}/{self.nb_epochs} train\",\n                leave=False,\n                total=len(self.train_data),\n            )\n            for step, (x, y) in loop:\n                x = x.to(self.device)\n                y = y.to(self.device)\n\n                # Train step\n                g_loss, d_loss = self.train_step(x, y)\n\n                train_g_loss += g_loss\n                train_d_loss += d_loss\n\n                loop.set_postfix_str(\n                    f\"g_loss: {train_g_loss / (step + 1) :.2f}, d_loss: {train_d_loss / (step + 1) :.2f}\"\n                )\n\n            # Validation loop\n            self.generator.eval()\n            self.discriminator.eval()\n\n            loop = tqdm(\n                enumerate(self.val_data),\n                desc=f\"Epoch {i + 1}/{self.nb_epochs} validation\",\n                leave=True,\n                total=len(self.val_data),\n            )\n            for step, (x, y) in loop:\n                x = x.to(self.device)\n                y = y.to(self.device)\n\n                # Validation step\n                g_loss, d_loss = self.val_step(x, y)\n\n                val_g_loss += g_loss\n                val_d_loss += d_loss\n\n                loop.set_postfix_str(\n                    f\"g_loss: {val_g_loss / (step + 1) :.2f} d_loss: {val_d_loss / (step + 1) :.2f}\"\n                )\n\n            # Saving best model\n            if self.save_path and best_score > val_g_loss:\n                best_score = val_g_loss\n                self.save_model()\n\n            # Log\n            self.log_metrics(\n                step=i,\n                train_g_loss=train_g_loss,\n                train_d_loss=train_d_loss,\n                val_g_loss=val_g_loss,\n                val_d_loss=val_d_loss,\n            )\n        return self.logs\n\n    def train_step(\n        self,\n        x: torch.Tensor,\n        y: torch.Tensor\n    ) -> tuple:\n        self.g_optimizer.zero_grad(set_to_none=True)\n        self.d_optimizer.zero_grad(set_to_none=True)\n        \n        fake_ = self.discriminator(\n            self.generator(x)\n        )\n        real = self.discriminator(y)\n        d_loss = (\n            torch.nn.functional.mse_loss(real, torch.ones_like(real, device=self.device)) + \n            torch.nn.functional.mse_loss(fake_, torch.zeros_like(fake_, device=self.device))\n        )\n\n        d_loss.backward()\n        self.d_optimizer.step()\n        \n        fake = self.discriminator(\n            self.generator(x)\n        )\n        \n        g_loss = torch.nn.functional.mse_loss(fake, torch.ones_like(fake, device=self.device))\n        g_loss.backward()\n        self.g_optimizer.step()\n\n        return g_loss.item(), d_loss.item()\n\n    @torch.no_grad()\n    def val_step(\n        self,\n        x: torch.Tensor,\n        y: torch.Tensor,\n    ) -> tuple:\n        fake = self.discriminator(\n            self.generator(x)\n        )\n        real = self.discriminator(y)\n        g_loss = torch.nn.functional.mse_loss(fake, torch.ones_like(fake, device=self.device))\n        d_loss = (\n            torch.nn.functional.mse_loss(real, torch.ones_like(real, device=self.device)) + \n            torch.nn.functional.mse_loss(fake, torch.zeros_like(fake, device=self.device))\n        )\n        return g_loss.item(), d_loss.item()\n\n    @torch.no_grad()\n    def log_metrics(\n        self,\n        step: int,\n        train_g_loss: torch.Tensor,\n        train_d_loss: torch.Tensor,\n        val_g_loss: torch.Tensor,\n        val_d_loss: torch.Tensor,\n    ):\n        self.logs[\"Step\"].append(step)\n        self.logs[\"Train_g_loss\"].append(train_g_loss / len(self.train_data))\n        self.logs[\"Train_d_loss\"].append(train_d_loss / len(self.train_data))\n        self.logs[\"Val_g_loss\"].append(val_g_loss / len(self.val_data))\n        self.logs[\"Val_d_loss\"].append(val_d_loss / len(self.val_data))\n        self.logs[\"Samples\"].append(make_grid(self.generator(self.z).cpu() * 0.5 + 0.5, normalize=True))\n\n    def save_model(self, full: bool = False):\n        if full:\n            torch.save(self.generator, Path(self.save_path) / \"generator.pth\")\n            torch.save(self.discriminator, Path(self.save_path) / \"discriminator.pth\")\n        else:\n            torch.save(\n                self.generator.state_dict(),\n                Path(self.save_path) / \"generator_weights.pth\",\n            )\n            torch.save(\n                self.discriminator.state_dict(),\n                Path(self.save_path) / \"discriminator_weights.pth\",\n            )\n\n    def load_model(self, full: bool = False):\n        if (\n            full\n            and (self.save_path / \"generator.pth\").is_file()\n            and (self.save_path / \"discriminator.pth\").is_file()\n        ):\n            self.generator = torch.load(self.save_path / \"generator.pth\")\n            self.discriminator = torch.load(self.save_path / \"discriminator.pth\")\n        elif (\n            self.save_path / \"generator_weights.pth\"\n        ).is_file() and (\n            self.save_path / \"discriminator_weights.pth\"\n        ).is_file():\n            self.generator.load_state_dict(torch.load(self.save_path / \"generator_weights.pth\"))\n            self.discriminator.load_state_dict(torch.load(self.save_path / \"discriminator_weights.pth\"))\n","metadata":{"execution":{"iopub.status.busy":"2022-09-22T15:11:36.117882Z","iopub.execute_input":"2022-09-22T15:11:36.118286Z","iopub.status.idle":"2022-09-22T15:11:36.148450Z","shell.execute_reply.started":"2022-09-22T15:11:36.118252Z","shell.execute_reply":"2022-09-22T15:11:36.147345Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Train the models","metadata":{}},{"cell_type":"code","source":"trainer = Trainer(\n        train_data = train_dataloader,\n        val_data = test_dataloader,\n        generator = Generator(in_channels=3, nb_layers=4, device=device),\n        discriminator = Discriminator(nb_layers=5, device=device),\n        nb_epochs = 30,\n        device = device,\n    )","metadata":{"execution":{"iopub.status.busy":"2022-09-22T15:11:45.527207Z","iopub.execute_input":"2022-09-22T15:11:45.527580Z","iopub.status.idle":"2022-09-22T15:11:46.479521Z","shell.execute_reply.started":"2022-09-22T15:11:45.527546Z","shell.execute_reply":"2022-09-22T15:11:46.478230Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"trainer.device","metadata":{"execution":{"iopub.status.busy":"2022-09-22T15:10:49.826592Z","iopub.execute_input":"2022-09-22T15:10:49.827055Z","iopub.status.idle":"2022-09-22T15:10:49.836776Z","shell.execute_reply.started":"2022-09-22T15:10:49.827016Z","shell.execute_reply":"2022-09-22T15:10:49.835764Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"trainer.init_optimizers()\nlogs = trainer.train()","metadata":{"execution":{"iopub.status.busy":"2022-09-22T13:08:01.711753Z","iopub.execute_input":"2022-09-22T13:08:01.712058Z","iopub.status.idle":"2022-09-22T13:09:59.873260Z","shell.execute_reply.started":"2022-09-22T13:08:01.712031Z","shell.execute_reply":"2022-09-22T13:09:59.871830Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Performance preview","metadata":{}},{"cell_type":"code","source":"fig, ax = plt.subplots(1, 2, figsize=(20, 4))\nax[0].plot(logs[\"Step\"], logs[\"Train_g_loss\"], label=\"Train_g_loss\")\nax[1].plot(logs[\"Step\"], logs[\"Train_d_loss\"], label=\"Train_d_loss\")\nax[0].plot(logs[\"Step\"], logs[\"Val_g_loss\"], label=\"Val_g_loss\")\nax[1].plot(logs[\"Step\"], logs[\"Val_d_loss\"], label=\"Val_d_loss\")\nax[0].set_title(\"Generator loss\")\nax[1].set_title(\"Discriminator loss\")\nax[0].legend()\nax[1].legend()\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2022-09-22T13:09:59.875722Z","iopub.execute_input":"2022-09-22T13:09:59.876360Z","iopub.status.idle":"2022-09-22T13:10:00.215816Z","shell.execute_reply.started":"2022-09-22T13:09:59.876317Z","shell.execute_reply":"2022-09-22T13:10:00.214899Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"fig, axs = plt.subplots(3, 3, figsize=(25, 15))\nj, k = 0, 0\nstep = (len(logs[\"Step\"])-1) // 9\nfor i in range(0, len(logs[\"Step\"]) - 1, step if step > 0 else 1):\n    axs[j, k].imshow(logs[\"Samples\"][i+step  if step > 0 else i].permute(1, 2, 0))\n    axs[j, k].set_title(f\"Generated image at step {logs['Step'][i+step  if step > 0 else i]}\")\n    axs[j, k].set(xticks=[], yticks=[])\n    if k >= 2:\n        j += 1\n        k = 0\n    else:\n        k += 1\n    if k+j == 4:\n        break\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2022-09-22T13:10:00.217404Z","iopub.execute_input":"2022-09-22T13:10:00.218058Z","iopub.status.idle":"2022-09-22T13:10:01.445622Z","shell.execute_reply.started":"2022-09-22T13:10:00.218017Z","shell.execute_reply":"2022-09-22T13:10:01.442682Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Create submission file","metadata":{}},{"cell_type":"code","source":"! mkdir ../images","metadata":{"execution":{"iopub.status.busy":"2022-09-22T15:11:49.341767Z","iopub.execute_input":"2022-09-22T15:11:49.342842Z","iopub.status.idle":"2022-09-22T15:11:50.322601Z","shell.execute_reply.started":"2022-09-22T15:11:49.342787Z","shell.execute_reply":"2022-09-22T15:11:50.321423Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"with torch.no_grad():\n    for i, (img, _) in enumerate(DataLoader(dataset, batch_size=1, shuffle=False)):\n        prediction = trainer.generator(img.to(device))[0].cpu().permute(1, 2, 0).numpy().astype(np.uint8)\n#         prediction = trainer.generator(torch.randn((1, 50), device=device))[0].cpu().permute(1, 2, 0).numpy().astype(np.uint8)\n        im = Image.fromarray(prediction)\n        im.save(f\"../images/{i}.jpg\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import shutil\nshutil.make_archive(\"/kaggle/working/images\", 'zip', \"/kaggle/images\")","metadata":{"execution":{"iopub.status.busy":"2022-09-22T12:47:30.286321Z","iopub.status.idle":"2022-09-22T12:47:30.287205Z","shell.execute_reply.started":"2022-09-22T12:47:30.286912Z","shell.execute_reply":"2022-09-22T12:47:30.286937Z"},"trusted":true},"execution_count":null,"outputs":[]}]}