{"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 warnings\nwarnings.filterwarnings(\"ignore\")","metadata":{"execution":{"iopub.status.busy":"2023-07-18T16:40:05.565262Z","iopub.execute_input":"2023-07-18T16:40:05.565935Z","iopub.status.idle":"2023-07-18T16:40:05.600417Z","shell.execute_reply.started":"2023-07-18T16:40:05.565894Z","shell.execute_reply":"2023-07-18T16:40:05.599203Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import glob\nimport os\nimport shutil","metadata":{"execution":{"iopub.status.busy":"2023-07-18T16:40:05.603474Z","iopub.execute_input":"2023-07-18T16:40:05.603938Z","iopub.status.idle":"2023-07-18T16:40:05.609634Z","shell.execute_reply.started":"2023-07-18T16:40:05.603897Z","shell.execute_reply":"2023-07-18T16:40:05.608429Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import matplotlib.pyplot as plt\nimport numpy as np\nimport torch\nimport torch.nn.functional as F\nimport torchvision.transforms as T\nfrom torch import nn\nfrom torch.utils.data import DataLoader, Dataset\nfrom torchvision.io import read_image\nfrom torchvision.utils import make_grid, save_image","metadata":{"execution":{"iopub.status.busy":"2023-07-18T16:40:05.625815Z","iopub.execute_input":"2023-07-18T16:40:05.626539Z","iopub.status.idle":"2023-07-18T16:40:10.096061Z","shell.execute_reply.started":"2023-07-18T16:40:05.626505Z","shell.execute_reply":"2023-07-18T16:40:10.094799Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"torch.manual_seed(0)\nnp.random.seed(0)","metadata":{"execution":{"iopub.status.busy":"2023-07-18T16:40:10.100807Z","iopub.execute_input":"2023-07-18T16:40:10.101471Z","iopub.status.idle":"2023-07-18T16:40:10.114113Z","shell.execute_reply.started":"2023-07-18T16:40:10.101433Z","shell.execute_reply":"2023-07-18T16:40:10.112850Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def show_img(img_tensor, nrow, title=\"\"):\n    \"\"\"\n    Display the given image tensor grid.\n\n    Args:\n        img_tensor (torch.Tensor): Image tensor grid.\n        nrow (int): Number of images per row.\n        title (str): Title for the plot (default: \"\").\n\n    \"\"\"\n    img_tensor = img_tensor.detach().cpu() * 0.5 + 0.5\n    img_grid = make_grid(img_tensor, nrow=nrow).permute(1, 2, 0)\n    plt.figure(figsize=(18, 8))\n    plt.imshow(img_grid)\n    plt.axis(\"off\")\n    plt.title(title)\n    plt.show()","metadata":{"execution":{"iopub.status.busy":"2023-07-18T16:40:10.115710Z","iopub.execute_input":"2023-07-18T16:40:10.116184Z","iopub.status.idle":"2023-07-18T16:40:10.126258Z","shell.execute_reply.started":"2023-07-18T16:40:10.116150Z","shell.execute_reply":"2023-07-18T16:40:10.125095Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CustomTransform(object):\n    \"\"\"\n    Custom data transformation class for image preprocessing.\n\n    Args:\n        load_dim (int): Size for resizing images during training (default: 286).\n        target_dim (int): Target size for cropping images during training (default: 256).\n\n    \"\"\"\n    def __init__(self, load_dim=286, target_dim=256):\n        self.transform_train = T.Compose([\n            T.Resize((load_dim, load_dim), interpolation=T.InterpolationMode.BICUBIC),\n            T.RandomCrop((target_dim, target_dim)),\n            T.RandomHorizontalFlip(),\n            T.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),\n        ])\n        self.transform = T.Resize((target_dim, target_dim), interpolation=T.InterpolationMode.BICUBIC)\n        \n    def __call__(self, img, stage):\n        if stage == \"fit\":\n            img = self.transform_train(img)\n        else:\n            img = self.transform(img)\n        return img * 2 - 1","metadata":{"execution":{"iopub.status.busy":"2023-07-18T16:40:10.129481Z","iopub.execute_input":"2023-07-18T16:40:10.129956Z","iopub.status.idle":"2023-07-18T16:40:10.142526Z","shell.execute_reply.started":"2023-07-18T16:40:10.129911Z","shell.execute_reply":"2023-07-18T16:40:10.141437Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CustomDataset(Dataset):\n    \"\"\"\n    Custom dataset class for loading and preprocessing image data.\n\n    Args:\n        filenames (list): List of file paths for the images.\n        transform (CustomTransform): Custom data transformation object.\n        stage (str): Stage of dataset (\"fit\" for training, None for validation/testing).\n\n    \"\"\"\n    def __init__(self, filenames, transform, stage):\n        self.filenames = filenames\n        self.transform = transform\n        self.stage = stage\n\n    def __len__(self):\n        return len(self.filenames)\n\n    def __getitem__(self, idx):\n        img_name = self.filenames[idx]\n        img = read_image(img_name) / 255.0\n        return self.transform(img, stage=self.stage)","metadata":{"execution":{"iopub.status.busy":"2023-07-18T16:40:10.143802Z","iopub.execute_input":"2023-07-18T16:40:10.144626Z","iopub.status.idle":"2023-07-18T16:40:10.156590Z","shell.execute_reply.started":"2023-07-18T16:40:10.144591Z","shell.execute_reply":"2023-07-18T16:40:10.155349Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"DM_CONFIG = {\n    \"monet_dir\": os.path.join(\"/kaggle/input/gan-getting-started/monet_jpg\", \"*.jpg\"),\n    \"photo_dir\": os.path.join(\"/kaggle/input/gan-getting-started/photo_jpg\", \"*.jpg\"),\n\n    \"loader_config\": {\n        \"num_workers\": os.cpu_count(),\n        \"pin_memory\": torch.cuda.is_available(),\n    },\n\n    \"sample_size\": 5,\n    \"batch_size\": 1,\n}","metadata":{"execution":{"iopub.status.busy":"2023-07-18T16:40:10.158923Z","iopub.execute_input":"2023-07-18T16:40:10.159497Z","iopub.status.idle":"2023-07-18T16:40:10.174656Z","shell.execute_reply.started":"2023-07-18T16:40:10.159455Z","shell.execute_reply":"2023-07-18T16:40:10.173389Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CustomDataModule:\n    \"\"\"\n    Custom data module class for managing data loading.\n\n    Args:\n        monet_dir (str): Directory path containing Monet paintings.\n        photo_dir (str): Directory path containing real photos.\n        loader_config (dict): Configuration for data loader.\n        sample_size (int): Number of samples for validation/testing.\n        batch_size (int): Batch size for training.\n\n    \"\"\"\n    def __init__(\n        self,\n        monet_dir,\n        photo_dir,\n        loader_config,\n        sample_size,\n        batch_size,\n    ):\n        self.loader_config = loader_config\n        self.sample_size = sample_size\n        self.batch_size = batch_size\n        self.monet_filenames = sorted(glob.glob(monet_dir))\n        self.photo_filenames = sorted(glob.glob(photo_dir))\n        self.transform = CustomTransform()\n\n    def train_dataloader(self):\n        loader_config = {\n            \"shuffle\": True,\n            \"drop_last\": True,\n            \"batch_size\": self.batch_size,\n            **self.loader_config,\n        }\n        dataset = CustomDataset(self.monet_filenames, self.transform, stage=\"fit\")\n        loader = DataLoader(dataset, **loader_config)\n        return loader\n\n    def val_dataloader(self):\n        loader_config = {\n            \"batch_size\": self.sample_size,\n            **self.loader_config,\n        }\n        dataset = CustomDataset(self.photo_filenames, self.transform, stage=None)\n        loader = DataLoader(dataset, **loader_config)\n        return loader\n\n    def test_dataloader(self):\n        return self.val_dataloader()","metadata":{"execution":{"iopub.status.busy":"2023-07-18T16:40:10.176560Z","iopub.execute_input":"2023-07-18T16:40:10.177311Z","iopub.status.idle":"2023-07-18T16:40:10.195711Z","shell.execute_reply.started":"2023-07-18T16:40:10.177272Z","shell.execute_reply":"2023-07-18T16:40:10.194318Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Downsampling(nn.Module):\n    \"\"\"\n    Downsampling block used in the generator and discriminator.\n\n    Args:\n        in_channels (int): Number of input channels.\n        out_channels (int): Number of output channels.\n        kernel_size (int): Kernel size for the convolutional layer (default: 4).\n        stride (int): Stride for the convolutional layer (default: 2).\n        padding (int): Padding for the convolutional layer (default: 1).\n        norm (bool): Whether to use instance normalization (default: True).\n        lrelu (bool): Whether to use leaky ReLU activation (default: True).\n\n    \"\"\"\n    def __init__(\n        self,\n        in_channels,\n        out_channels,\n        kernel_size=4,\n        stride=2,\n        padding=1,\n        norm=True,\n        lrelu=True,\n    ):\n        super().__init__()\n        layers = [\n            nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, padding=padding, bias=not norm),\n        ]\n        if norm:\n            layers.append(nn.InstanceNorm2d(out_channels, affine=True))\n        if lrelu is not None:\n            layers.append(nn.LeakyReLU(0.2, True) if lrelu else nn.ReLU(True))\n        self.block = nn.Sequential(*layers)\n\n    def forward(self, x):\n        return self.block(x)","metadata":{"execution":{"iopub.status.busy":"2023-07-18T16:40:10.197473Z","iopub.execute_input":"2023-07-18T16:40:10.197901Z","iopub.status.idle":"2023-07-18T16:40:10.212507Z","shell.execute_reply.started":"2023-07-18T16:40:10.197857Z","shell.execute_reply":"2023-07-18T16:40:10.211231Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Upsampling(nn.Module):\n    \"\"\"\n    Upsampling block used in the generator.\n\n    Args:\n        in_channels (int): Number of input channels.\n        out_channels (int): Number of output channels.\n        dropout (bool): Whether to use dropout regularization (default: False).\n\n    \"\"\"\n    def __init__(self, in_channels, out_channels, dropout=False):\n        super(Upsampling, self).__init__()\n        layers = [\n            nn.ConvTranspose2d(in_channels, out_channels, kernel_size=4, stride=2, padding=1, output_padding=0),\n            nn.InstanceNorm2d(out_channels),\n            nn.ReLU(inplace=True),\n        ]\n        if dropout:\n            layers.append(nn.Dropout2d(0.5))\n        self.block = nn.Sequential(*layers)\n\n    def forward(self, x):\n        return self.block(x)","metadata":{"execution":{"iopub.status.busy":"2023-07-18T16:40:10.214523Z","iopub.execute_input":"2023-07-18T16:40:10.215040Z","iopub.status.idle":"2023-07-18T16:40:10.228263Z","shell.execute_reply.started":"2023-07-18T16:40:10.214974Z","shell.execute_reply":"2023-07-18T16:40:10.227304Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class ReflectionPad2d(nn.Module):\n    \"\"\"\n    Reflection padding layer used in the generator.\n\n    \"\"\"\n    def forward(self, x):\n        _, _, h, w = x.size()\n        padding = (1, 1, 1, 1)  # Adjust padding size based on input tensor shape\n        return nn.functional.pad(x, padding, mode='reflect')","metadata":{"execution":{"iopub.status.busy":"2023-07-18T16:40:10.233682Z","iopub.execute_input":"2023-07-18T16:40:10.235550Z","iopub.status.idle":"2023-07-18T16:40:10.245076Z","shell.execute_reply.started":"2023-07-18T16:40:10.235496Z","shell.execute_reply":"2023-07-18T16:40:10.243709Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class ResBlock(nn.Module):\n    \"\"\"\n    Residual block used in the generator.\n\n    Args:\n        channels (int): Number of input and output channels.\n\n    \"\"\"\n    def __init__(self, channels):\n        super(ResBlock, self).__init__()\n        self.channels = channels\n        self.block = nn.Sequential(\n            nn.Conv2d(channels, channels, kernel_size=3, stride=1, padding=1, bias=False),\n            nn.ReLU(inplace=True),\n            nn.Conv2d(channels, channels, kernel_size=3, stride=1, padding=1, bias=False),\n            nn.ReLU(inplace=True)\n        )\n        self.downsample = nn.Conv2d(channels, channels, kernel_size=1, stride=1, padding=0, bias=False)\n        \n    def forward(self, x):\n        residual = x\n        x = self.block(x)\n        \n        # Skip instance normalization if spatial dimensions are 1x1\n        if x.shape[2:] != (1, 1):\n            x = F.instance_norm(x)\n        \n        return residual + x","metadata":{"execution":{"iopub.status.busy":"2023-07-18T16:40:10.246662Z","iopub.execute_input":"2023-07-18T16:40:10.247271Z","iopub.status.idle":"2023-07-18T16:40:10.263566Z","shell.execute_reply.started":"2023-07-18T16:40:10.247224Z","shell.execute_reply":"2023-07-18T16:40:10.261884Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Generator(nn.Module):\n    \"\"\"\n    Generator model for image-to-image translation.\n\n    Args:\n        hid_channels (int): Number of channels in the hidden layers.\n        in_channels (int): Number of input channels.\n        out_channels (int): Number of output channels.\n        num_resblocks (int): Number of residual blocks.\n\n    \"\"\"\n    def __init__(self, hid_channels, in_channels, out_channels, num_resblocks):\n        super().__init__()\n        self.downsampling_path = nn.Sequential(\n            Downsampling(in_channels, hid_channels, norm=False),  # 64x128x128\n            Downsampling(hid_channels, hid_channels*2),  # 128x64x64\n            Downsampling(hid_channels*2, hid_channels*4),  # 256x32x32\n            Downsampling(hid_channels*4, hid_channels*8),  # 512x16x16\n            Downsampling(hid_channels*8, hid_channels*8),  # 512x8x8\n            Downsampling(hid_channels*8, hid_channels*8),  # 512x4x4\n            Downsampling(hid_channels*8, hid_channels*8),  # 512x2x2\n            Downsampling(hid_channels*8, hid_channels*8, norm=False),  # 512x1x1, instance norm does not work on 1x1\n        )\n        self.upsampling_path = nn.Sequential(\n            Upsampling(hid_channels*8, hid_channels*8, dropout=True),  # (512+512)x2x2\n            Upsampling(hid_channels*16, hid_channels*8, dropout=True),  # (512+512)x4x4\n            Upsampling(hid_channels*16, hid_channels*8, dropout=True),  # (512+512)x8x8\n            Upsampling(hid_channels*16, hid_channels*8),  # (512+512)x16x16\n            Upsampling(hid_channels*16, hid_channels*4),  # (256+256)x32x32\n            Upsampling(hid_channels*8, hid_channels*2),  # (128+128)x64x64\n            Upsampling(hid_channels*4, hid_channels),  # (64+64)x128x128\n        )\n        self.feature_block = nn.Sequential(\n            nn.ConvTranspose2d(hid_channels*2, out_channels,\n                               kernel_size=4, stride=2, padding=1),  # 3x256x256\n            nn.Tanh(),\n        )\n        \n    def forward(self, x):\n        skips = []\n        for down in self.downsampling_path:\n            x = down(x)\n            skips.append(x)\n        skips = reversed(skips[:-1])\n\n        for up, skip in zip(self.upsampling_path, skips):\n            x = up(x)\n            x = torch.cat([x, skip], dim=1)\n        return self.feature_block(x)","metadata":{"execution":{"iopub.status.busy":"2023-07-18T16:40:10.266258Z","iopub.execute_input":"2023-07-18T16:40:10.266857Z","iopub.status.idle":"2023-07-18T16:40:10.290555Z","shell.execute_reply.started":"2023-07-18T16:40:10.266796Z","shell.execute_reply":"2023-07-18T16:40:10.289046Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Discriminator(nn.Module):\n    \"\"\"\n    Discriminator model for distinguishing real and generated images.\n\n    Args:\n        hid_channels (int): Number of channels in the hidden layers.\n        in_channels (int): Number of input channels.\n\n    \"\"\"\n    def __init__(self, hid_channels, in_channels=3):\n        super().__init__()\n        self.block = nn.Sequential(\n            Downsampling(in_channels, hid_channels, norm=False),  # 64x128x128\n            Downsampling(hid_channels, hid_channels*2),  # 128x64x64\n            Downsampling(hid_channels*2, hid_channels*4),  # 256x32x32\n            Downsampling(hid_channels*4, hid_channels*8, stride=1),  # 512x31x31\n            nn.Conv2d(hid_channels*8, 1, kernel_size=4, padding=1),  # 1x30x30\n        )\n        \n    def forward(self, x):\n        return self.block(x)","metadata":{"execution":{"iopub.status.busy":"2023-07-18T16:40:10.292182Z","iopub.execute_input":"2023-07-18T16:40:10.292620Z","iopub.status.idle":"2023-07-18T16:40:10.303948Z","shell.execute_reply.started":"2023-07-18T16:40:10.292571Z","shell.execute_reply":"2023-07-18T16:40:10.302971Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"hid_channels = 64\nin_channels = 3\nout_channels = 3\nnum_resblocks = 9","metadata":{"execution":{"iopub.status.busy":"2023-07-18T16:40:10.305663Z","iopub.execute_input":"2023-07-18T16:40:10.306055Z","iopub.status.idle":"2023-07-18T16:40:10.316087Z","shell.execute_reply.started":"2023-07-18T16:40:10.305994Z","shell.execute_reply":"2023-07-18T16:40:10.314798Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"gen = Generator(hid_channels, in_channels, out_channels, num_resblocks)\ndisc = Discriminator(in_channels)","metadata":{"execution":{"iopub.status.busy":"2023-07-18T16:40:10.317921Z","iopub.execute_input":"2023-07-18T16:40:10.318405Z","iopub.status.idle":"2023-07-18T16:40:11.163031Z","shell.execute_reply.started":"2023-07-18T16:40:10.318364Z","shell.execute_reply":"2023-07-18T16:40:11.161680Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def init_weights(m):\n    \"\"\"\n    Initialize the weights of convolutional layers and instance normalization layers.\n\n    Args:\n        m (nn.Module): Module to initialize.\n\n    \"\"\"\n    if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d)):\n        nn.init.normal_(m.weight.data, 0.0, 0.02)\n        if m.bias is not None:\n            nn.init.constant_(m.bias.data, 0.0)\n    elif isinstance(m, nn.InstanceNorm2d):\n        if m.weight is not None:\n            nn.init.normal_(m.weight.data, 1.0, 0.02)\n        if m.bias is not None:\n            nn.init.constant_(m.bias.data, 0.0)","metadata":{"execution":{"iopub.status.busy":"2023-07-18T16:40:11.164736Z","iopub.execute_input":"2023-07-18T16:40:11.165200Z","iopub.status.idle":"2023-07-18T16:40:11.174656Z","shell.execute_reply.started":"2023-07-18T16:40:11.165155Z","shell.execute_reply":"2023-07-18T16:40:11.173092Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"gen.apply(init_weights)\ndisc.apply(init_weights)","metadata":{"execution":{"iopub.status.busy":"2023-07-18T16:40:11.177922Z","iopub.execute_input":"2023-07-18T16:40:11.179735Z","iopub.status.idle":"2023-07-18T16:40:11.670228Z","shell.execute_reply.started":"2023-07-18T16:40:11.179693Z","shell.execute_reply":"2023-07-18T16:40:11.669269Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\ngen.to(device)\ndisc.to(device)","metadata":{"execution":{"iopub.status.busy":"2023-07-18T16:40:11.671173Z","iopub.execute_input":"2023-07-18T16:40:11.671489Z","iopub.status.idle":"2023-07-18T16:40:11.694870Z","shell.execute_reply.started":"2023-07-18T16:40:11.671463Z","shell.execute_reply":"2023-07-18T16:40:11.693707Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"gen_optim = torch.optim.Adam(gen.parameters(), lr=2e-4, betas=(0.5, 0.999))\ndisc_optim = torch.optim.Adam(disc.parameters(), lr=2e-4, betas=(0.5, 0.999))","metadata":{"execution":{"iopub.status.busy":"2023-07-18T16:40:11.696960Z","iopub.execute_input":"2023-07-18T16:40:11.697938Z","iopub.status.idle":"2023-07-18T16:40:11.710284Z","shell.execute_reply.started":"2023-07-18T16:40:11.697886Z","shell.execute_reply":"2023-07-18T16:40:11.708487Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"adv_criterion = nn.BCEWithLogitsLoss()\nrecon_criterion = nn.L1Loss()","metadata":{"execution":{"iopub.status.busy":"2023-07-18T16:40:11.712355Z","iopub.execute_input":"2023-07-18T16:40:11.712852Z","iopub.status.idle":"2023-07-18T16:40:11.726523Z","shell.execute_reply.started":"2023-07-18T16:40:11.712804Z","shell.execute_reply":"2023-07-18T16:40:11.725443Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def train_step(gen, disc, gen_optim, disc_optim, adv_criterion, recon_criterion, real, monet):\n    \"\"\"\n    Perform a single training step for the generator and discriminator models.\n\n    Args:\n        gen (nn.Module): Generator model.\n        disc (nn.Module): Discriminator model.\n        gen_optim (torch.optim.Optimizer): Optimizer for the generator.\n        disc_optim (torch.optim.Optimizer): Optimizer for the discriminator.\n        adv_criterion (nn.Module): Adversarial loss criterion.\n        recon_criterion (nn.Module): Reconstruction loss criterion.\n        real (torch.Tensor): Real images.\n        monet (torch.Tensor): Monet paintings.\n\n    Returns:\n        float: Generator loss.\n        float: Discriminator loss.\n\n    \"\"\"\n    gen_optim.zero_grad()\n    disc_optim.zero_grad()\n\n    real = real.to(device)\n    monet = monet.to(device)\n\n    # Train Generator\n    fake_monet = gen(real)\n    disc_fake_monet = disc(fake_monet)\n    gen_adv_loss = adv_criterion(disc_fake_monet, torch.ones_like(disc_fake_monet))\n    gen_recon_loss = recon_criterion(fake_monet, monet)\n    gen_loss = gen_adv_loss + 10 * gen_recon_loss\n\n    gen_loss.backward()\n    gen_optim.step()\n\n    # Train Discriminator\n    disc_optim.zero_grad()\n\n    disc_real_monet = disc(monet)\n    disc_fake_monet = disc(fake_monet.detach())\n    disc_real_loss = adv_criterion(disc_real_monet, torch.ones_like(disc_real_monet))\n    disc_fake_loss = adv_criterion(disc_fake_monet, torch.zeros_like(disc_fake_monet))\n    disc_loss = 0.5 * (disc_real_loss + disc_fake_loss)\n\n    disc_loss.backward()\n    disc_optim.step()\n\n    return gen_loss.item(), disc_loss.item()","metadata":{"execution":{"iopub.status.busy":"2023-07-18T16:40:11.728437Z","iopub.execute_input":"2023-07-18T16:40:11.729255Z","iopub.status.idle":"2023-07-18T16:40:11.743801Z","shell.execute_reply.started":"2023-07-18T16:40:11.729210Z","shell.execute_reply":"2023-07-18T16:40:11.742415Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def train(gen, disc, gen_optim, disc_optim, adv_criterion, recon_criterion, train_loader, val_loader, num_epochs):\n    \"\"\"\n    Train the generator and discriminator models.\n\n    Args:\n        gen (nn.Module): Generator model.\n        disc (nn.Module): Discriminator model.\n        gen_optim (torch.optim.Optimizer): Optimizer for the generator.\n        disc_optim (torch.optim.Optimizer): Optimizer for the discriminator.\n        adv_criterion (nn.Module): Adversarial loss criterion.\n        recon_criterion (nn.Module): Reconstruction loss criterion.\n        train_loader (torch.utils.data.DataLoader): DataLoader for training data.\n        val_loader (torch.utils.data.DataLoader): DataLoader for validation data.\n        num_epochs (int): Number of training epochs.\n\n    Returns:\n        list: Generator losses per epoch.\n        list: Discriminator losses per epoch.\n\n    \"\"\"\n    gen_losses = []\n    disc_losses = []\n\n    for epoch in range(num_epochs):\n        gen_loss = 0.0\n        disc_loss = 0.0\n        for real in train_loader:\n            monet = next(iter(train_loader))\n\n            cur_gen_loss, cur_disc_loss = train_step(gen, disc, gen_optim, disc_optim, adv_criterion, recon_criterion, real, monet)\n            gen_loss += cur_gen_loss\n            disc_loss += cur_disc_loss\n\n        gen_losses.append(gen_loss / len(train_loader))\n        disc_losses.append(disc_loss / len(train_loader))\n\n        print(f\"Epoch [{epoch+1}/{num_epochs}], Gen Loss: {gen_losses[-1]:.4f}, Disc Loss: {disc_losses[-1]:.4f}\")\n\n        if (epoch+1) % 10 == 0:\n            val_imgs = next(iter(val_loader))\n            val_imgs = val_imgs.to(device)\n            with torch.no_grad():\n                gen.eval()\n                gen_imgs = gen(val_imgs)\n                gen.train()\n\n            show_img(val_imgs, nrow=5, title=f\"Epoch [{epoch+1}/{num_epochs}], Validation Real Images\")\n            show_img(gen_imgs, nrow=5, title=f\"Epoch [{epoch+1}/{num_epochs}], Validation Generated Images\")\n\n    return gen_losses, disc_losses","metadata":{"execution":{"iopub.status.busy":"2023-07-18T16:40:11.746203Z","iopub.execute_input":"2023-07-18T16:40:11.747219Z","iopub.status.idle":"2023-07-18T16:40:11.763107Z","shell.execute_reply.started":"2023-07-18T16:40:11.747172Z","shell.execute_reply":"2023-07-18T16:40:11.761682Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"dm = CustomDataModule(**DM_CONFIG)\ntrain_loader = dm.train_dataloader()  # Added line to define train_loader","metadata":{"execution":{"iopub.status.busy":"2023-07-18T16:40:11.765499Z","iopub.execute_input":"2023-07-18T16:40:11.766076Z","iopub.status.idle":"2023-07-18T16:40:12.204687Z","shell.execute_reply.started":"2023-07-18T16:40:11.766029Z","shell.execute_reply":"2023-07-18T16:40:12.203550Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"num_epochs = 10\nval_dir = \"/kaggle/input/gan-getting-started/photo_jpg\"\nval_filenames = sorted(glob.glob(os.path.join(val_dir, \"*.jpg\")))\ntransform = CustomTransform(load_dim=286, target_dim=256)\nval_dataset = CustomDataset(val_filenames, transform, stage=\"fit\")\nbatch_size = 1\nnum_workers = 0\nval_loader = torch.utils.data.DataLoader(val_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers)\ngen_losses, disc_losses = train(gen, disc, gen_optim, disc_optim, adv_criterion, recon_criterion, train_loader, val_loader, num_epochs)","metadata":{"execution":{"iopub.status.busy":"2023-07-18T16:40:12.207032Z","iopub.execute_input":"2023-07-18T16:40:12.208227Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plt.figure(figsize=(5, 5))\nplt.plot(gen_losses, label=\"Generator Loss\")\nplt.plot(disc_losses, label=\"Discriminator Loss\")\nplt.xlabel(\"Epochs\")\nplt.ylabel(\"Loss\")\nplt.title(\"Training Losses\")\nplt.legend()\nplt.show()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def save_checkpoint(model, optimizer, filename):\n    \"\"\"\n    Save model checkpoint to a file.\n\n    Args:\n        model (nn.Module): Model to save.\n        optimizer (torch.optim.Optimizer): Optimizer associated with the model.\n        filename (str): File path to save the checkpoint.\n\n    \"\"\"\n    checkpoint = {\n        \"model_state_dict\": model.state_dict(),\n        \"optimizer_state_dict\": optimizer.state_dict(),\n    }\n    torch.save(checkpoint, filename)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"save_checkpoint(gen, gen_optim, \"generator.pth\")\nsave_checkpoint(disc, disc_optim, \"discriminator.pth\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def generate_images(model, val_loader, filename):\n    \"\"\"\n    Generate images using the trained model and save them.\n\n    Args:\n        model (nn.Module): Trained model.\n        val_loader (torch.utils.data.DataLoader): DataLoader for validation data.\n        filename (str): File path to save the generated images.\n\n    \"\"\"\n    model.eval()\n    with torch.no_grad():\n        val_imgs = next(iter(val_loader))\n        val_imgs = val_imgs.to(device)\n        gen_imgs = model(val_imgs)\n\n    save_image(gen_imgs, filename, normalize=True, nrow=5)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"generate_images(gen, val_loader, \"images.png\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"shutil.make_archive(\"images\", \"zip\", \".\", \"images.png\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}