{"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\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\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":"2022-10-10T06:31:45.158718Z","iopub.execute_input":"2022-10-10T06:31:45.159661Z","iopub.status.idle":"2022-10-10T06:31:51.189091Z","shell.execute_reply.started":"2022-10-10T06:31:45.159531Z","shell.execute_reply":"2022-10-10T06:31:51.188133Z"},"collapsed":true,"jupyter":{"outputs_hidden":true},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nimport numpy as np\nimport matplotlib.pyplot as plt\nimport torch\nimport torch.nn as nn\nimport torchvision\nfrom torch.utils.data import Dataset, DataLoader, random_split\nfrom torchvision.utils import make_grid\nfrom PIL import Image\n\nfrom tqdm import tqdm\nimport shutil\n\nfrom pathlib import Path","metadata":{"execution":{"iopub.status.busy":"2022-10-10T06:31:51.190635Z","iopub.execute_input":"2022-10-10T06:31:51.191200Z","iopub.status.idle":"2022-10-10T06:31:53.268873Z","shell.execute_reply.started":"2022-10-10T06:31:51.191168Z","shell.execute_reply":"2022-10-10T06:31:53.267880Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"data_root = '/kaggle/input/gan-getting-started/'\nprint(data_root)\nprint(os.listdir(data_root))\nmonet_path = 'monet_jpg'\nphoto_path = 'photo_jpg'","metadata":{"execution":{"iopub.status.busy":"2022-10-10T06:31:53.270505Z","iopub.execute_input":"2022-10-10T06:31:53.271142Z","iopub.status.idle":"2022-10-10T06:31:53.277201Z","shell.execute_reply.started":"2022-10-10T06:31:53.271103Z","shell.execute_reply":"2022-10-10T06:31:53.276253Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class MonetPhoto(Dataset):\n    def __init__(self, data_root, monet_path, photo_path, transform=None):\n        self.data_root = data_root\n        self.monet_path = monet_path\n        self.monet_images = os.listdir(os.path.join(data_root, monet_path))\n        self.photo_path = photo_path\n        self.photo_images = os.listdir(os.path.join(data_root, photo_path))\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(os.path.join(data_root, monet_path, self.monet_images[idx % len(self.monet_images)]))\n        photo_image = Image.open(os.path.join(data_root, photo_path, 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-10-10T06:31:53.280116Z","iopub.execute_input":"2022-10-10T06:31:53.280815Z","iopub.status.idle":"2022-10-10T06:31:53.290639Z","shell.execute_reply.started":"2022-10-10T06:31:53.280777Z","shell.execute_reply":"2022-10-10T06:31:53.289377Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"transforms = torchvision.transforms.Compose([\n    torchvision.transforms.Resize((256,256)),\n    torchvision.transforms.ToTensor(),\n    torchvision.transforms.Normalize(0.5, 0.5)\n])\n\ndataset = MonetPhoto(data_root, monet_path, photo_path, transform=transforms)","metadata":{"execution":{"iopub.status.busy":"2022-10-10T06:31:53.292548Z","iopub.execute_input":"2022-10-10T06:31:53.293163Z","iopub.status.idle":"2022-10-10T06:31:53.307352Z","shell.execute_reply.started":"2022-10-10T06:31:53.293124Z","shell.execute_reply":"2022-10-10T06:31:53.306426Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"len(dataset)","metadata":{"execution":{"iopub.status.busy":"2022-10-10T06:31:53.308966Z","iopub.execute_input":"2022-10-10T06:31:53.309416Z","iopub.status.idle":"2022-10-10T06:31:53.320041Z","shell.execute_reply.started":"2022-10-10T06:31:53.309382Z","shell.execute_reply":"2022-10-10T06:31:53.319042Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"x, y = dataset.__getitem__(1)\nx.shape, y.shape","metadata":{"execution":{"iopub.status.busy":"2022-10-10T06:31:53.323195Z","iopub.execute_input":"2022-10-10T06:31:53.324483Z","iopub.status.idle":"2022-10-10T06:31:53.359907Z","shell.execute_reply.started":"2022-10-10T06:31:53.324456Z","shell.execute_reply":"2022-10-10T06:31:53.358829Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"x.max(), x.min(), y.max(), y.min()","metadata":{"execution":{"iopub.status.busy":"2022-10-10T06:31:53.361210Z","iopub.execute_input":"2022-10-10T06:31:53.361586Z","iopub.status.idle":"2022-10-10T06:31:53.371402Z","shell.execute_reply.started":"2022-10-10T06:31:53.361551Z","shell.execute_reply":"2022-10-10T06:31:53.370219Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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-10-10T06:31:53.373103Z","iopub.execute_input":"2022-10-10T06:31:53.373722Z","iopub.status.idle":"2022-10-10T06:31:53.379443Z","shell.execute_reply.started":"2022-10-10T06:31:53.373666Z","shell.execute_reply":"2022-10-10T06:31:53.378196Z"},"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-10-10T06:31:53.384722Z","iopub.execute_input":"2022-10-10T06:31:53.385635Z","iopub.status.idle":"2022-10-10T06:31:53.391706Z","shell.execute_reply.started":"2022-10-10T06:31:53.385598Z","shell.execute_reply":"2022-10-10T06:31:53.390573Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"len(training_data)","metadata":{"execution":{"iopub.status.busy":"2022-10-10T06:31:53.393436Z","iopub.execute_input":"2022-10-10T06:31:53.393941Z","iopub.status.idle":"2022-10-10T06:31:53.402737Z","shell.execute_reply.started":"2022-10-10T06:31:53.393855Z","shell.execute_reply":"2022-10-10T06:31:53.401498Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_dataloader = DataLoader(training_data.dataset, batch_size=32, shuffle=True, num_workers=2)\ntest_dataloader = DataLoader(testing_data.dataset, batch_size=32, shuffle=True, num_workers=2)","metadata":{"execution":{"iopub.status.busy":"2022-10-10T06:31:53.403962Z","iopub.execute_input":"2022-10-10T06:31:53.404490Z","iopub.status.idle":"2022-10-10T06:31:53.412698Z","shell.execute_reply.started":"2022-10-10T06:31:53.404456Z","shell.execute_reply":"2022-10-10T06:31:53.411659Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"mones, photos = next(iter(train_dataloader))\nplt.imshow(mones[0].permute(1, 2, 0).numpy() * 0.5 + 0.5)","metadata":{"execution":{"iopub.status.busy":"2022-10-10T06:31:53.416213Z","iopub.execute_input":"2022-10-10T06:31:53.416601Z","iopub.status.idle":"2022-10-10T06:31:54.933640Z","shell.execute_reply.started":"2022-10-10T06:31:53.416544Z","shell.execute_reply":"2022-10-10T06:31:54.932484Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plt.imshow(photos[0].permute(1, 2, 0).numpy() * 0.5 + 0.5)","metadata":{"execution":{"iopub.status.busy":"2022-10-10T06:31:54.935257Z","iopub.execute_input":"2022-10-10T06:31:54.937989Z","iopub.status.idle":"2022-10-10T06:31:55.173314Z","shell.execute_reply.started":"2022-10-10T06:31:54.937945Z","shell.execute_reply":"2022-10-10T06:31:55.172447Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\ndevice","metadata":{"execution":{"iopub.status.busy":"2022-10-10T06:31:55.174806Z","iopub.execute_input":"2022-10-10T06:31:55.175416Z","iopub.status.idle":"2022-10-10T06:31:55.244615Z","shell.execute_reply.started":"2022-10-10T06:31:55.175380Z","shell.execute_reply":"2022-10-10T06:31:55.243398Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CNNBlock(nn.Module):\n    def __init__(self, in_channels, out_channels, stride=2, device=\"cuda\"):\n        super().__init__()\n        self.device = device\n        self.conv = nn.Sequential(\n            nn.Conv2d(in_channels, out_channels, 4, stride, bias=False, padding_mode='reflect', device=self.device),\n            nn.BatchNorm2d(out_channels, device=self.device),\n            nn.LeakyReLU(0.2),\n        )\n        \n    def forward(self, x):\n        return self.conv(x)","metadata":{"execution":{"iopub.status.busy":"2022-10-10T06:31:55.246945Z","iopub.execute_input":"2022-10-10T06:31:55.247828Z","iopub.status.idle":"2022-10-10T06:31:55.256435Z","shell.execute_reply.started":"2022-10-10T06:31:55.247787Z","shell.execute_reply":"2022-10-10T06:31:55.255363Z"},"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], device=\"cuda\"):\n        super().__init__()\n        self.device = device\n        self.initial = nn.Sequential(\n            nn.Conv2d(in_channels, features[0], kernel_size=4, stride=2, padding=1, padding_mode='reflect', device=self.device),\n            nn.LeakyReLU(0.2)\n        )\n        \n        layers = []\n        in_channels = features[0]\n        for feature in features[1:]:\n            layers.append(\n                CNNBlock(in_channels, feature, stride=1 if feature == features[-1] else 2, device=self.device),\n            )\n            in_channels = feature\n        \n        layers.append(\n            nn.Conv2d(\n                in_channels, 1, kernel_size=4, stride=1, padding=1, padding_mode='reflect', device=self.device\n            )\n        )\n            \n        self.model = nn.Sequential(*layers)\n        \n    def forward(self, x):\n        x = self.initial(x)\n        return self.model(x)","metadata":{"execution":{"iopub.status.busy":"2022-10-10T06:31:55.258389Z","iopub.execute_input":"2022-10-10T06:31:55.258757Z","iopub.status.idle":"2022-10-10T06:31:55.271947Z","shell.execute_reply.started":"2022-10-10T06:31:55.258728Z","shell.execute_reply":"2022-10-10T06:31:55.271334Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def test():\n    x = (torch.randn((1,3, 256,256))).to('cuda')\n    model = Discriminator().to('cuda')\n    preds = model(x)\n    print(preds.shape)\n    \ntest()","metadata":{"execution":{"iopub.status.busy":"2022-10-10T06:31:55.273218Z","iopub.execute_input":"2022-10-10T06:31:55.275526Z","iopub.status.idle":"2022-10-10T06:32:03.801089Z","shell.execute_reply.started":"2022-10-10T06:31:55.275480Z","shell.execute_reply":"2022-10-10T06:32:03.799919Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Block(nn.Module):\n    def __init__(self, in_channels, out_channels, down=True, act = 'relu', use_dropout=False, device=\"cuda\"):\n        super().__init__()\n        self.device = device\n        self.conv = nn.Sequential(\n            nn.Conv2d(in_channels, out_channels, 4, 2, 1, bias=False, padding_mode='reflect', device=self.device)\n            if down\n            else nn.ConvTranspose2d(in_channels, out_channels, 4, 2, 1, bias=False, device=self.device),\n            nn.BatchNorm2d(out_channels, device=self.device),\n            nn.ReLU() if act=='relu' else nn.LeakyReLU(0.2),\n        )\n        self.use_dropout = use_dropout\n        self.dropout = nn.Dropout(0.5)\n    \n    def forward(self, x):\n        x = self.conv(x)\n        return self.dropout(x) if self.use_dropout else x","metadata":{"execution":{"iopub.status.busy":"2022-10-10T06:32:03.802588Z","iopub.execute_input":"2022-10-10T06:32:03.803241Z","iopub.status.idle":"2022-10-10T06:32:03.813388Z","shell.execute_reply.started":"2022-10-10T06:32:03.803201Z","shell.execute_reply":"2022-10-10T06:32:03.811964Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Generator(nn.Module):\n    def __init__(self, in_channels=3, features=64, device=\"cuda\"):\n        super().__init__()\n        self.device = device\n        self.initial_down = nn.Sequential(\n            nn.Conv2d(in_channels, features, 4, 2, 1, padding_mode='reflect', device=self.device),\n            nn.LeakyReLU(0.2),\n        ) #128\n        \n        self.down1 = Block(features, features*2, down=True, act='leaky', use_dropout=False)   #64\n        self.down2 = Block(features*2, features*4, down=True, act='leaky', use_dropout=False) #32\n        self.down3 = Block(features*4, features*8, down=True, act='leaky', use_dropout=False) #16\n        self.down4 = Block(features*8, features*8, down=True, act='leaky', use_dropout=False) #8\n        self.down5 = Block(features*8, features*8, down=True, act='leaky', use_dropout=False) #4\n        self.down6 = Block(features*8, features*8, down=True, act='leaky', use_dropout=False) #2\n        self.bottleneck = nn.Sequential(\n            nn.Conv2d(features*8, features*8, 4, 2, 1, padding_mode='reflect', device=self.device), #1x1\n            nn.ReLU(),\n        )\n        self.up1 = Block(features*8, features*8, down=False, act='relu', use_dropout=True)\n        self.up2 = Block(features*8*2, features*8, down=False, act='relu', use_dropout=True)\n        self.up3 = Block(features*8*2, features*8, down=False, act='relu', use_dropout=True)\n        self.up4 = Block(features*8*2, features*8, down=False, act='relu', use_dropout=False)\n        self.up5 = Block(features*8*2, features*4, down=False, act='relu', use_dropout=False)\n        self.up6 = Block(features*4*2, features*2, down=False, act='relu', use_dropout=False)\n        self.up7 = Block(features*2*2, features, down=False, act='relu', use_dropout=False)\n        self.finil_up = nn.Sequential(\n            nn.ConvTranspose2d(features*2, in_channels, kernel_size=4, stride=2, padding=1, device=self.device),\n            nn.Tanh(),\n        )\n        \n    def forward(self, x):\n        d1 = self.initial_down(x)\n        d2 = self.down1(d1)\n        d3 = self.down2(d2)\n        d4 = self.down3(d3)\n        d5 = self.down4(d4)\n        d6 = self.down5(d5)\n        d7 = self.down6(d6)\n        \n        bottleneck = self.bottleneck(d7)\n        \n        up1 = self.up1(bottleneck)\n        up2 = self.up2(torch.cat([up1, d7], 1))\n        up3 = self.up3(torch.cat([up2, d6], 1))\n        up4 = self.up4(torch.cat([up3, d5], 1))\n        up5 = self.up5(torch.cat([up4, d4], 1))\n        up6 = self.up6(torch.cat([up5, d3], 1))\n        up7 = self.up7(torch.cat([up6, d2], 1))\n        \n        return self.finil_up(torch.cat([up7, d1], 1))","metadata":{"execution":{"iopub.status.busy":"2022-10-10T06:32:03.815328Z","iopub.execute_input":"2022-10-10T06:32:03.816400Z","iopub.status.idle":"2022-10-10T06:32:03.837091Z","shell.execute_reply.started":"2022-10-10T06:32:03.816221Z","shell.execute_reply":"2022-10-10T06:32:03.836051Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def test():\n    x = torch.randn((1,3, 256,256)).to('cuda')\n    model = Generator()\n    preds = model(x)\n    print(preds.shape)\n    \ntest()","metadata":{"execution":{"iopub.status.busy":"2022-10-10T06:32:03.838656Z","iopub.execute_input":"2022-10-10T06:32:03.839292Z","iopub.status.idle":"2022-10-10T06:32:03.900978Z","shell.execute_reply.started":"2022-10-10T06:32:03.839247Z","shell.execute_reply":"2022-10-10T06:32:03.899818Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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=3e-4, 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            # plot\n            fake_img = self.generator((photos[0].unsqueeze(0)).to('cuda'))\n            fake_img = fake_img[0].cpu().detach()\n            \n            fig, ax = plt.subplots(1, 2, figsize=(20, 4))\n            ax[0].imshow(photos[0].permute(1, 2, 0).numpy() * 0.5 + 0.5)\n            ax[0].set_title(\"Photo\")\n            ax[0].set(xticks=[], yticks=[])\n            ax[1].imshow(fake_img.permute(1, 2, 0).numpy() * 0.5 + 0.5  )\n            ax[1].set_title(\"FakeMonet\")\n            ax[1].set(xticks=[], yticks=[])\n            plt.show()\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\"))","metadata":{"execution":{"iopub.status.busy":"2022-10-10T06:32:03.902670Z","iopub.execute_input":"2022-10-10T06:32:03.903218Z","iopub.status.idle":"2022-10-10T06:32:03.936752Z","shell.execute_reply.started":"2022-10-10T06:32:03.903180Z","shell.execute_reply":"2022-10-10T06:32:03.935664Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!mkdir models","metadata":{"execution":{"iopub.status.busy":"2022-10-10T06:32:03.938357Z","iopub.execute_input":"2022-10-10T06:32:03.939258Z","iopub.status.idle":"2022-10-10T06:32:04.965257Z","shell.execute_reply.started":"2022-10-10T06:32:03.939217Z","shell.execute_reply":"2022-10-10T06:32:04.963952Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"trainer = Trainer(\n        train_data = train_dataloader,\n        val_data = test_dataloader,\n        generator = Generator(in_channels=3, features=64),\n        discriminator = Discriminator(),\n        nb_epochs = 30,\n        device = device,\n        save_path = './models/'\n    )","metadata":{"execution":{"iopub.status.busy":"2022-10-10T06:32:04.968224Z","iopub.execute_input":"2022-10-10T06:32:04.968910Z","iopub.status.idle":"2022-10-10T06:32:06.290636Z","shell.execute_reply.started":"2022-10-10T06:32:04.968868Z","shell.execute_reply":"2022-10-10T06:32:06.289203Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"trainer.load_model()","metadata":{"execution":{"iopub.status.busy":"2022-10-10T06:32:06.292453Z","iopub.execute_input":"2022-10-10T06:32:06.292916Z","iopub.status.idle":"2022-10-10T06:32:06.299196Z","shell.execute_reply.started":"2022-10-10T06:32:06.292875Z","shell.execute_reply":"2022-10-10T06:32:06.297896Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"trainer.init_optimizers()\nlogs = trainer.train()","metadata":{"execution":{"iopub.status.busy":"2022-10-10T06:32:06.301098Z","iopub.execute_input":"2022-10-10T06:32:06.302071Z","iopub.status.idle":"2022-10-10T06:35:46.836650Z","shell.execute_reply.started":"2022-10-10T06:32:06.302015Z","shell.execute_reply":"2022-10-10T06:35:46.835606Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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-10-10T06:35:46.838525Z","iopub.execute_input":"2022-10-10T06:35:46.839298Z","iopub.status.idle":"2022-10-10T06:35:47.162398Z","shell.execute_reply.started":"2022-10-10T06:35:46.839256Z","shell.execute_reply":"2022-10-10T06:35:47.161479Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plt.imshow(photos[1].permute(1, 2, 0).numpy() * 0.5 + 0.5)","metadata":{"execution":{"iopub.status.busy":"2022-10-10T06:35:47.168015Z","iopub.execute_input":"2022-10-10T06:35:47.168347Z","iopub.status.idle":"2022-10-10T06:35:47.408432Z","shell.execute_reply.started":"2022-10-10T06:35:47.168319Z","shell.execute_reply":"2022-10-10T06:35:47.407464Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"fake_img = trainer.generator((photos[1].unsqueeze(0)).to('cuda'))\nfake_img = fake_img[0].cpu().detach()\nplt.imshow(fake_img.permute(1, 2, 0).numpy() * 0.5 + 0.5  )","metadata":{"execution":{"iopub.status.busy":"2022-10-10T06:35:47.409812Z","iopub.execute_input":"2022-10-10T06:35:47.410912Z","iopub.status.idle":"2022-10-10T06:35:47.647415Z","shell.execute_reply.started":"2022-10-10T06:35:47.410873Z","shell.execute_reply":"2022-10-10T06:35:47.646506Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"! mkdir ../images","metadata":{"execution":{"iopub.status.busy":"2022-10-10T06:35:47.648794Z","iopub.execute_input":"2022-10-10T06:35:47.649754Z","iopub.status.idle":"2022-10-10T06:35:48.640995Z","shell.execute_reply.started":"2022-10-10T06:35:47.649718Z","shell.execute_reply":"2022-10-10T06:35:48.639661Z"},"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()* 0.5 + 0.5\n        im = Image.fromarray((prediction * 255).astype(np.uint8))\n        im.save(\"../images/\" + str(i) + \".jpg\")","metadata":{"execution":{"iopub.status.busy":"2022-10-10T06:35:48.642973Z","iopub.execute_input":"2022-10-10T06:35:48.643597Z","iopub.status.idle":"2022-10-10T06:37:48.475101Z","shell.execute_reply.started":"2022-10-10T06:35:48.643527Z","shell.execute_reply":"2022-10-10T06:37:48.474065Z"},"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-10-10T06:37:48.476488Z","iopub.execute_input":"2022-10-10T06:37:48.476935Z","iopub.status.idle":"2022-10-10T06:37:57.607352Z","shell.execute_reply.started":"2022-10-10T06:37:48.476892Z","shell.execute_reply":"2022-10-10T06:37:57.606372Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}