{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":21755,"databundleVersionId":1475600,"sourceType":"competition"}],"dockerImageVersionId":30823,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"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","trusted":true,"execution":{"iopub.status.busy":"2024-12-21T04:04:52.845080Z","iopub.execute_input":"2024-12-21T04:04:52.845371Z","iopub.status.idle":"2024-12-21T04:05:02.884437Z","shell.execute_reply.started":"2024-12-21T04:04:52.845341Z","shell.execute_reply":"2024-12-21T04:05:02.883709Z"}},"outputs":[],"execution_count":null},{"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":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-21T04:05:02.885701Z","iopub.execute_input":"2024-12-21T04:05:02.886120Z","iopub.status.idle":"2024-12-21T04:05:07.689398Z","shell.execute_reply.started":"2024-12-21T04:05:02.886096Z","shell.execute_reply":"2024-12-21T04:05:07.688540Z"}},"outputs":[],"execution_count":null},{"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":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-21T04:05:07.690525Z","iopub.execute_input":"2024-12-21T04:05:07.690886Z","iopub.status.idle":"2024-12-21T04:05:07.695621Z","shell.execute_reply.started":"2024-12-21T04:05:07.690864Z","shell.execute_reply":"2024-12-21T04:05:07.694792Z"}},"outputs":[],"execution_count":null},{"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":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-21T04:05:07.696464Z","iopub.execute_input":"2024-12-21T04:05:07.696691Z","iopub.status.idle":"2024-12-21T04:05:07.712406Z","shell.execute_reply.started":"2024-12-21T04:05:07.696660Z","shell.execute_reply":"2024-12-21T04:05:07.711771Z"}},"outputs":[],"execution_count":null},{"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":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-21T04:05:07.713138Z","iopub.execute_input":"2024-12-21T04:05:07.713330Z","iopub.status.idle":"2024-12-21T04:05:07.735750Z","shell.execute_reply.started":"2024-12-21T04:05:07.713313Z","shell.execute_reply":"2024-12-21T04:05:07.735115Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"len(dataset)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-21T04:05:07.737660Z","iopub.execute_input":"2024-12-21T04:05:07.737852Z","iopub.status.idle":"2024-12-21T04:05:07.748087Z","shell.execute_reply.started":"2024-12-21T04:05:07.737835Z","shell.execute_reply":"2024-12-21T04:05:07.747352Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"x, y = dataset.__getitem__(1)\nx.shape, y.shape","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-21T04:05:07.749284Z","iopub.execute_input":"2024-12-21T04:05:07.749567Z","iopub.status.idle":"2024-12-21T04:05:07.871417Z","shell.execute_reply.started":"2024-12-21T04:05:07.749539Z","shell.execute_reply":"2024-12-21T04:05:07.870678Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"x.max(), x.min(), y.max(), y.min()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-21T04:05:07.872104Z","iopub.execute_input":"2024-12-21T04:05:07.872320Z","iopub.status.idle":"2024-12-21T04:05:07.908343Z","shell.execute_reply.started":"2024-12-21T04:05:07.872301Z","shell.execute_reply":"2024-12-21T04:05:07.907688Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_size, test_size = int(len(dataset) * 0.8), len(dataset) - int(len(dataset) * 0.8)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-21T04:05:07.909174Z","iopub.execute_input":"2024-12-21T04:05:07.909373Z","iopub.status.idle":"2024-12-21T04:05:07.912988Z","shell.execute_reply.started":"2024-12-21T04:05:07.909355Z","shell.execute_reply":"2024-12-21T04:05:07.912166Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"training_data, testing_data = random_split(dataset, [train_size, test_size])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-21T04:05:07.913788Z","iopub.execute_input":"2024-12-21T04:05:07.914110Z","iopub.status.idle":"2024-12-21T04:05:07.934055Z","shell.execute_reply.started":"2024-12-21T04:05:07.914079Z","shell.execute_reply":"2024-12-21T04:05:07.933408Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"len(training_data)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-21T04:05:07.934851Z","iopub.execute_input":"2024-12-21T04:05:07.935133Z","iopub.status.idle":"2024-12-21T04:05:07.942240Z","shell.execute_reply.started":"2024-12-21T04:05:07.935105Z","shell.execute_reply":"2024-12-21T04:05:07.941425Z"}},"outputs":[],"execution_count":null},{"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":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-21T04:05:07.942949Z","iopub.execute_input":"2024-12-21T04:05:07.943236Z","iopub.status.idle":"2024-12-21T04:05:07.959168Z","shell.execute_reply.started":"2024-12-21T04:05:07.943198Z","shell.execute_reply":"2024-12-21T04:05:07.958389Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"mones, photos = next(iter(train_dataloader))\nplt.imshow(mones[0].permute(1, 2, 0).numpy() * 0.5 + 0.5)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-21T04:05:07.960025Z","iopub.execute_input":"2024-12-21T04:05:07.960309Z","iopub.status.idle":"2024-12-21T04:05:09.504660Z","shell.execute_reply.started":"2024-12-21T04:05:07.960282Z","shell.execute_reply":"2024-12-21T04:05:09.503694Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"plt.imshow(photos[0].permute(1, 2, 0).numpy() * 0.5 + 0.5)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-21T04:05:09.505646Z","iopub.execute_input":"2024-12-21T04:05:09.505889Z","iopub.status.idle":"2024-12-21T04:05:09.803461Z","shell.execute_reply.started":"2024-12-21T04:05:09.505863Z","shell.execute_reply":"2024-12-21T04:05:09.802579Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nprint(\"Using device:\", device)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-21T04:05:09.804279Z","iopub.execute_input":"2024-12-21T04:05:09.804492Z","iopub.status.idle":"2024-12-21T04:05:09.895562Z","shell.execute_reply.started":"2024-12-21T04:05:09.804464Z","shell.execute_reply":"2024-12-21T04:05:09.894696Z"}},"outputs":[],"execution_count":null},{"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":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-21T04:05:09.896494Z","iopub.execute_input":"2024-12-21T04:05:09.896830Z","iopub.status.idle":"2024-12-21T04:05:09.909271Z","shell.execute_reply.started":"2024-12-21T04:05:09.896799Z","shell.execute_reply":"2024-12-21T04:05:09.908618Z"}},"outputs":[],"execution_count":null},{"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":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-21T04:05:09.910064Z","iopub.execute_input":"2024-12-21T04:05:09.910284Z","iopub.status.idle":"2024-12-21T04:05:09.927089Z","shell.execute_reply.started":"2024-12-21T04:05:09.910265Z","shell.execute_reply":"2024-12-21T04:05:09.926249Z"}},"outputs":[],"execution_count":null},{"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":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-21T04:05:09.930258Z","iopub.execute_input":"2024-12-21T04:05:09.930481Z","iopub.status.idle":"2024-12-21T04:05:10.953578Z","shell.execute_reply.started":"2024-12-21T04:05:09.930461Z","shell.execute_reply":"2024-12-21T04:05:10.952685Z"}},"outputs":[],"execution_count":null},{"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":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-21T04:05:10.955174Z","iopub.execute_input":"2024-12-21T04:05:10.955471Z","iopub.status.idle":"2024-12-21T04:05:10.960888Z","shell.execute_reply.started":"2024-12-21T04:05:10.955447Z","shell.execute_reply":"2024-12-21T04:05:10.959889Z"}},"outputs":[],"execution_count":null},{"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":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-21T04:05:10.961664Z","iopub.execute_input":"2024-12-21T04:05:10.961912Z","iopub.status.idle":"2024-12-21T04:05:10.981965Z","shell.execute_reply.started":"2024-12-21T04:05:10.961892Z","shell.execute_reply":"2024-12-21T04:05:10.980981Z"}},"outputs":[],"execution_count":null},{"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":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-21T04:05:10.982747Z","iopub.execute_input":"2024-12-21T04:05:10.982943Z","iopub.status.idle":"2024-12-21T04:05:11.225165Z","shell.execute_reply.started":"2024-12-21T04:05:10.982926Z","shell.execute_reply":"2024-12-21T04:05:11.224425Z"}},"outputs":[],"execution_count":null},{"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":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-21T04:05:11.226127Z","iopub.execute_input":"2024-12-21T04:05:11.226364Z","iopub.status.idle":"2024-12-21T04:05:11.245723Z","shell.execute_reply.started":"2024-12-21T04:05:11.226344Z","shell.execute_reply":"2024-12-21T04:05:11.244972Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!mkdir models","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-21T04:05:11.246491Z","iopub.execute_input":"2024-12-21T04:05:11.246802Z","iopub.status.idle":"2024-12-21T04:05:11.395628Z","shell.execute_reply.started":"2024-12-21T04:05:11.246773Z","shell.execute_reply":"2024-12-21T04:05:11.394709Z"}},"outputs":[],"execution_count":null},{"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":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-21T04:05:11.396704Z","iopub.execute_input":"2024-12-21T04:05:11.396926Z","iopub.status.idle":"2024-12-21T04:05:12.386940Z","shell.execute_reply.started":"2024-12-21T04:05:11.396905Z","shell.execute_reply":"2024-12-21T04:05:12.385780Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"trainer.load_model()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-21T04:05:12.388153Z","iopub.execute_input":"2024-12-21T04:05:12.388494Z","iopub.status.idle":"2024-12-21T04:05:12.392703Z","shell.execute_reply.started":"2024-12-21T04:05:12.388459Z","shell.execute_reply":"2024-12-21T04:05:12.391873Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"trainer.init_optimizers()\nlogs = trainer.train()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-21T04:05:12.393535Z","iopub.execute_input":"2024-12-21T04:05:12.393816Z"}},"outputs":[],"execution_count":null},{"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":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"plt.imshow(photos[1].permute(1, 2, 0).numpy() * 0.5 + 0.5)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"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":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"! mkdir ../images","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"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":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import shutil\nshutil.make_archive(\"/kaggle/working/images\", 'zip', \"/kaggle/images\")","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}