Add training code

BIG_GAN.ipynb

@@ -0,0 +1,581 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch \n",
+    "import torch.nn as nn\n",
+    "import torch.nn.functional as F\n",
+    "\n",
+    "def orthogonal_regularization(weight):\n",
+    "    weight = weight.flatten(1)\n",
+    "    return torch.norm(\n",
+    "        torch.dot(weight, weight) * (torch.ones_like(weight) - torch.eye(weight.shape[0]))\n",
+    "    )\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class ClassConditionalBatchNorm2d(nn.Module):\n",
+    "    def __init__(self, in_channels, out_channels):\n",
+    "        super().__init__()\n",
+    "        self.bn = nn.BatchNorm2d(out_channels)\n",
+    "        self.class_scale_transform = nn.utils.spectral_norm(nn.Linear(in_channels, out_channels))\n",
+    "        self.class_shift_transform = nn.utils.spectral_norm(nn.Linear(in_channels, out_channels))\n",
+    "\n",
+    "    def forward(self, x, y):\n",
+    "        normalized_image = self.bn(x)\n",
+    "        class_scale = (1 + self.class_scale_transform(y))[:,:,None,None]\n",
+    "        class_shift = self.class_shift_transform(y)[:,:,None,None]\n",
+    "        tranformed_image = class_scale * normalized_image + class_shift\n",
+    "        return tranformed_image\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Attention Block for Big GAN\n",
+    "This paper has all that you need to implement this class:\n",
+    "https://arxiv.org/pdf/1805.08318"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class AttentionBlock(nn.Module):\n",
+    "    '''\n",
+    "    AttentionBlock Class\n",
+    "    Values:\n",
+    "    channels: number of channels in input\n",
+    "    '''\n",
+    "    def __init__(self, channels):\n",
+    "        super().__init__()\n",
+    "\n",
+    "        self.channels = channels\n",
+    "\n",
+    "        self.theta = nn.utils.spectral_norm(nn.Conv2d(channels, channels // 8, kernel_size=1, padding=0, bias=False))\n",
+    "        self.phi = nn.utils.spectral_norm(nn.Conv2d(channels, channels // 8, kernel_size=1, padding=0, bias=False))\n",
+    "        self.g = nn.utils.spectral_norm(nn.Conv2d(channels, channels // 2, kernel_size=1, padding=0, bias=False))\n",
+    "        self.o = nn.utils.spectral_norm(nn.Conv2d(channels // 2, channels, kernel_size=1, padding=0, bias=False))\n",
+    "\n",
+    "        self.gamma = nn.Parameter(torch.tensor(0.), requires_grad=True)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        spatial_size = x.shape[2] * x.shape[3]\n",
+    "\n",
+    "        # Apply convolutions to get query (theta), key (phi), and value (g) transforms\n",
+    "        theta = self.theta(x)\n",
+    "        phi = F.max_pool2d(self.phi(x), kernel_size=2)\n",
+    "        g = F.max_pool2d(self.g(x), kernel_size=2)\n",
+    "\n",
+    "        # Reshape spatial size for self-attention\n",
+    "        theta = theta.view(-1, self.channels // 8, spatial_size)\n",
+    "        phi = phi.view(-1, self.channels // 8, spatial_size // 4)\n",
+    "        g = g.view(-1, self.channels // 2, spatial_size // 4)\n",
+    "\n",
+    "        # Compute dot product attention with query (theta) and key (phi) matrices\n",
+    "        beta = F.softmax(torch.bmm(theta.transpose(1, 2), phi), dim=-1)\n",
+    "\n",
+    "        # Compute scaled dot product attention with value (g) and attention (beta) matrices\n",
+    "        o = self.o(torch.bmm(g, beta.transpose(1, 2)).view(-1, self.channels // 2, x.shape[2], x.shape[3]))\n",
+    "\n",
+    "        # Apply gain and residual\n",
+    "        return self.gamma * o + x"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class GResidualBlock(nn.Module):\n",
+    "    '''\n",
+    "    GResidualBlock Class\n",
+    "    Values:\n",
+    "    c_dim: the dimension of conditional vector [c, z], a scalar\n",
+    "    in_channels: the number of channels in the input, a scalar\n",
+    "    out_channels: the number of channels in the output, a scalar\n",
+    "    '''\n",
+    "\n",
+    "    def __init__(self, c_dim, in_channels, out_channels):\n",
+    "        super().__init__()\n",
+    "\n",
+    "        self.conv1 = nn.utils.spectral_norm(nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1))\n",
+    "        self.conv2 = nn.utils.spectral_norm(nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1))\n",
+    "\n",
+    "        self.bn1 = ClassConditionalBatchNorm2d(c_dim, in_channels)\n",
+    "        self.bn2 = ClassConditionalBatchNorm2d(c_dim, out_channels)\n",
+    "\n",
+    "        self.activation = nn.ReLU()\n",
+    "        self.upsample_fn = nn.Upsample(scale_factor=2)     # upsample occurs in every gblock\n",
+    "\n",
+    "        self.mixin = (in_channels != out_channels)\n",
+    "        if self.mixin:\n",
+    "            self.conv_mixin = nn.utils.spectral_norm(nn.Conv2d(in_channels, out_channels, kernel_size=1, padding=0))\n",
+    "\n",
+    "    def forward(self, x, y):\n",
+    "        # h := upsample(x, y)\n",
+    "        h = self.bn1(x, y)\n",
+    "        h = self.activation(h)\n",
+    "        h = self.upsample_fn(h)\n",
+    "        h = self.conv1(h)\n",
+    "\n",
+    "        # h := conv(h, y)\n",
+    "        h = self.bn2(h, y)\n",
+    "        h = self.activation(h)\n",
+    "        h = self.conv2(h)\n",
+    "\n",
+    "        # x := upsample(x)\n",
+    "        x = self.upsample_fn(x)\n",
+    "        if self.mixin:\n",
+    "            x = self.conv_mixin(x)\n",
+    "\n",
+    "        return h + x"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class Generator(nn.Module):\n",
+    "    '''\n",
+    "    Generator Class\n",
+    "    Values:\n",
+    "    z_dim: the dimension of random noise sampled, a scalar\n",
+    "    shared_dim: the dimension of shared class embeddings, a scalar\n",
+    "    base_channels: the number of base channels, a scalar\n",
+    "    bottom_width: the height/width of image before it gets upsampled, a scalar\n",
+    "    n_classes: the number of image classes, a scalar\n",
+    "    '''\n",
+    " \n",
+    "    def __init__(self, base_channels=96, bottom_width=4, z_dim=120, shared_dim=128, n_classes=1000):\n",
+    "        super().__init__()\n",
+    " \n",
+    "        n_chunks = 6    # 5 (generator blocks) + 1 (generator input)\n",
+    "        self.z_chunk_size = z_dim // n_chunks\n",
+    "        self.z_dim = z_dim\n",
+    "        self.shared_dim = shared_dim\n",
+    "        self.bottom_width = bottom_width\n",
+    " \n",
+    "        # No spectral normalization on embeddings, which authors observe to cripple the generator\n",
+    "        self.shared_emb = nn.Embedding(n_classes, shared_dim)\n",
+    " \n",
+    "        self.proj_z = nn.Linear(self.z_chunk_size, 16 * base_channels * bottom_width ** 2)\n",
+    " \n",
+    "        # Can't use one big nn.Sequential since we are adding class+noise at each block\n",
+    "        self.g_blocks = nn.ModuleList([\n",
+    "            nn.ModuleList([\n",
+    "                GResidualBlock(shared_dim + self.z_chunk_size, 16 * base_channels, 16 * base_channels),\n",
+    "                AttentionBlock(16 * base_channels),\n",
+    "            ]),\n",
+    "            nn.ModuleList([\n",
+    "                GResidualBlock(shared_dim + self.z_chunk_size, 16 * base_channels, 8 * base_channels),\n",
+    "                AttentionBlock(8 * base_channels),\n",
+    "            ]),\n",
+    "            nn.ModuleList([\n",
+    "                GResidualBlock(shared_dim + self.z_chunk_size, 8 * base_channels, 4 * base_channels),\n",
+    "                AttentionBlock(4 * base_channels),\n",
+    "            ]),\n",
+    "            nn.ModuleList([\n",
+    "                GResidualBlock(shared_dim + self.z_chunk_size, 4 * base_channels, 2 * base_channels),\n",
+    "                AttentionBlock(2 * base_channels),\n",
+    "            ]),\n",
+    "            nn.ModuleList([\n",
+    "                GResidualBlock(shared_dim + self.z_chunk_size, 2 * base_channels, base_channels),\n",
+    "                AttentionBlock(base_channels),\n",
+    "            ]),\n",
+    "        ])\n",
+    "        self.proj_o = nn.Sequential(\n",
+    "            nn.BatchNorm2d(base_channels),\n",
+    "            nn.ReLU(inplace=True),\n",
+    "            nn.utils.spectral_norm(nn.Conv2d(base_channels, 3, kernel_size=1, padding=0)),\n",
+    "            nn.Tanh(),\n",
+    "        )\n",
+    " \n",
+    "    def forward(self, z, y):\n",
+    "        '''\n",
+    "        z: random noise with size self.z_dim\n",
+    "        y: class embeddings with size self.shared_dim\n",
+    "            = NOTE =\n",
+    "            y should be class embeddings from self.shared_emb, not the raw class labels\n",
+    "        '''\n",
+    "        # Chunk z and concatenate to shared class embeddings\n",
+    "        zs = torch.split(z, self.z_chunk_size, dim=1)\n",
+    "        z = zs[0]\n",
+    "        ys = [torch.cat([y, z], dim=1) for z in zs[1:]]\n",
+    " \n",
+    "        # Project noise and reshape to feed through generator blocks\n",
+    "        h = self.proj_z(z)\n",
+    "        h = h.view(h.size(0), -1, self.bottom_width, self.bottom_width)\n",
+    " \n",
+    "        # Feed through generator blocks\n",
+    "        for idx, g_block in enumerate(self.g_blocks):\n",
+    "            h = g_block[0](h, ys[idx])\n",
+    "            h = g_block[1](h)\n",
+    " \n",
+    "        # Project to 3 RGB channels with tanh to map values to [-1, 1]\n",
+    "        h = self.proj_o(h)\n",
+    " \n",
+    "        return h"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class DResidualBlock(nn.Module):\n",
+    "    '''\n",
+    "    DResidualBlock Class\n",
+    "    Values:\n",
+    "    in_channels: the number of channels in the input, a scalar\n",
+    "    out_channels: the number of channels in the output, a scalar\n",
+    "    downsample: whether to apply downsampling\n",
+    "    use_preactivation: whether to apply an activation function before the first convolution\n",
+    "    '''\n",
+    "\n",
+    "    def __init__(self, in_channels, out_channels, downsample=True, use_preactivation=False):\n",
+    "        super().__init__()\n",
+    "\n",
+    "        self.conv1 = nn.utils.spectral_norm(nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1))\n",
+    "        self.conv2 = nn.utils.spectral_norm(nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1))\n",
+    "\n",
+    "        self.activation = nn.ReLU()\n",
+    "        self.use_preactivation = use_preactivation  # apply preactivation in all except first dblock\n",
+    "\n",
+    "        self.downsample = downsample    # downsample occurs in all except last dblock\n",
+    "        if downsample:\n",
+    "            self.downsample_fn = nn.AvgPool2d(2)\n",
+    "        self.mixin = (in_channels != out_channels) or downsample\n",
+    "        if self.mixin:\n",
+    "            self.conv_mixin = nn.utils.spectral_norm(nn.Conv2d(in_channels, out_channels, kernel_size=1, padding=0))\n",
+    "\n",
+    "    def _residual(self, x):\n",
+    "        if self.use_preactivation:\n",
+    "            if self.mixin:\n",
+    "                x = self.conv_mixin(x)\n",
+    "            if self.downsample:\n",
+    "                x = self.downsample_fn(x)\n",
+    "        else:\n",
+    "            if self.downsample:\n",
+    "                x = self.downsample_fn(x)\n",
+    "            if self.mixin:\n",
+    "                x = self.conv_mixin(x)\n",
+    "        return x\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        # Apply preactivation if applicable\n",
+    "        if self.use_preactivation:\n",
+    "            h = F.relu(x)\n",
+    "        else:\n",
+    "            h = x\n",
+    "\n",
+    "        h = self.conv1(h)\n",
+    "        h = self.activation(h)\n",
+    "        if self.downsample:\n",
+    "            h = self.downsample_fn(h)\n",
+    "\n",
+    "        return h + self._residual(x)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class Discriminator(nn.Module):\n",
+    "    def __init__(self, base_channels=96, n_classes = 38):\n",
+    "        super().__init__()\n",
+    "\n",
+    "        # Adding the class information\n",
+    "        self.shared_emb = nn.utils.spectral_norm(nn.Embedding(n_classes, 16 * base_channels))\n",
+    "\n",
+    "        self.d_blocks = nn.Sequential(\n",
+    "            DResidualBlock(3, base_channels, downsample=True, use_preactivation=False),\n",
+    "            AttentionBlock(base_channels),\n",
+    "\n",
+    "            DResidualBlock(base_channels, 2 * base_channels, downsample=True, use_preactivation=True),\n",
+    "            AttentionBlock(2 * base_channels),\n",
+    "\n",
+    "            DResidualBlock(2 * base_channels, 4 * base_channels, downsample=True, use_preactivation=True),\n",
+    "            AttentionBlock(4 * base_channels),\n",
+    "\n",
+    "            DResidualBlock(4 * base_channels, 8 * base_channels, downsample=True, use_preactivation=True),\n",
+    "            AttentionBlock(8 * base_channels),\n",
+    "\n",
+    "            DResidualBlock(8 * base_channels, 16 * base_channels, downsample=True, use_preactivation=True),\n",
+    "            AttentionBlock(16 * base_channels),\n",
+    "\n",
+    "            DResidualBlock(16 * base_channels, 16 * base_channels, downsample=False, use_preactivation=True),\n",
+    "            AttentionBlock(16 * base_channels),\n",
+    "\n",
+    "            nn.ReLU(inplace=True)\n",
+    "        )\n",
+    "\n",
+    "        self.proj_o = nn.utils.spectral_norm(nn.Linear(16 * base_channels, 1))\n",
+    "\n",
+    "    def forward(self, x, y=None):\n",
+    "        h = self.d_blocks(x)\n",
+    "        h = torch.sum(h, dim=[2,3])\n",
+    "\n",
+    "        # Class unconditional output\n",
+    "        uncond_out = self.proj_o(h)\n",
+    "        if y is None:\n",
+    "            return uncond_out\n",
+    "        \n",
+    "        # Class conditional output\n",
+    "        cond_out = torch.sum(self.shared_emb(y) * h, dim=1, keepdim=True)\n",
+    "        return uncond_out + cond_out\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "device = 'cuda'\n",
+    "\n",
+    "# Initialize models\n",
+    "base_channels = 96\n",
+    "z_dim = 120\n",
+    "n_classes = 38  \n",
+    "shared_dim = 128\n",
+    "generator = Generator(base_channels=base_channels, bottom_width=4, z_dim=z_dim, shared_dim=shared_dim, n_classes=n_classes).to(device)\n",
+    "discriminator = Discriminator(base_channels=base_channels, n_classes=n_classes).to(device)\n",
+    "\n",
+    "# Initialize weights orthogonally\n",
+    "for module in generator.modules():\n",
+    "    if (isinstance(module, nn.Conv2d) or isinstance(module, nn.Linear) or isinstance(module, nn.Embedding)):\n",
+    "        nn.init.orthogonal_(module.weight)\n",
+    "for module in discriminator.modules():\n",
+    "    if (isinstance(module, nn.Conv2d) or isinstance(module, nn.Linear) or isinstance(module, nn.Embedding)):\n",
+    "        nn.init.orthogonal_(module.weight)\n",
+    "\n",
+    "# Initialize optimizers\n",
+    "g_optimizer = torch.optim.Adam(generator.parameters(), lr=1e-4, betas=(0.0, 0.999), eps=1e-6)\n",
+    "d_optimizer = torch.optim.Adam(discriminator.parameters(), lr=4e-4, betas=(0.0, 0.999), eps=1e-6)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from datasets import load_dataset\n",
+    "dataset = load_dataset('YashNagraj75/gochiUSA')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from torch.utils.data import DataLoader\n",
+    "from torchvision import transforms\n",
+    "\n",
+    "\n",
+    "train_dataset = dataset['train']\n",
+    "test_dataset = dataset['test']\n",
+    "# Define transformations for both train and test datasets\n",
+    "transform = transforms.Compose([\n",
+    "    transforms.Resize((128, 128)),  # Resize to 128x128, adjust as needed\n",
+    "    transforms.ToTensor(),  # Convert to tensor\n",
+    "    transforms.Normalize((0.5,), (0.5,)),\n",
+    "])\n",
+    "\n",
+    "def transform_examples(example):\n",
+    "    example['image'] = transform(example['image'])\n",
+    "    return example\n",
+    "\n",
+    "transformed_train_dataset = train_dataset.map(transform_examples)\n",
+    "transformed_test_dataset = test_dataset.map(transform_examples)\n",
+    "\n",
+    "transformed_train_dataset.set_format('torch', columns=['image', 'label'])\n",
+    "transformed_test_dataset.set_format('torch', columns=['image', 'label'])\n",
+    "# DataLoader for training and testing\n",
+    "train_dataloader = DataLoader(transformed_train_dataset, batch_size=38, shuffle=True) \n",
+    "test_dataloader = DataLoader(transformed_test_dataset, batch_size=38, shuffle=False)   \n",
+    "# Example usage: iterate over the train DataLoader\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "torch.Tensor"
+      ]
+     },
+     "execution_count": 108,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "type(train_dataloader.dataset[0]['image'])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [
+    {
+     "ename": "NameError",
+     "evalue": "name 'train_dataloader' is not defined",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
+      "Cell \u001b[0;32mIn[1], line 17\u001b[0m\n\u001b[1;32m     13\u001b[0m display_step \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m500\u001b[39m\n\u001b[1;32m     16\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m epoch \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(num_epochs):\n\u001b[0;32m---> 17\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m i, batch \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28menumerate\u001b[39m(\u001b[43mtrain_dataloader\u001b[49m):\n\u001b[1;32m     18\u001b[0m         g_losses \u001b[38;5;241m=\u001b[39m []\n\u001b[1;32m     19\u001b[0m         d_losses \u001b[38;5;241m=\u001b[39m []\n",
+      "\u001b[0;31mNameError\u001b[0m: name 'train_dataloader' is not defined"
+     ]
+    }
+   ],
+   "source": [
+    "import torchvision.utils as vutils\n",
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "def critic_loss(real_output, fake_output):\n",
+    "    return torch.mean(fake_output) - torch.mean(real_output)\n",
+    "\n",
+    "def generator_loss(fake_output):\n",
+    "    return -torch.mean(fake_output)\n",
+    "\n",
+    "num_epochs = 200\n",
+    "batch_size = 38\n",
+    "n_critic = 5\n",
+    "display_step = 500\n",
+    "\n",
+    "\n",
+    "g_losses = []\n",
+    "d_losses = []\n",
+    "for epoch in range(num_epochs):\n",
+    "    for i, batch in enumerate(train_dataloader):\n",
+    "        real_images = batch[\"image\"].to(device)\n",
+    "        real_labels = batch[\"label\"].to(device)\n",
+    "\n",
+    "        \n",
+    "        d_optimizer.zero_grad()\n",
+    "        \n",
+    "        real_output = discriminator(real_images, real_labels)\n",
+    "        \n",
+    "        z = torch.randn(batch_size, z_dim, device=device)\n",
+    "        y = torch.arange(start=0, end=n_classes, device=device).long()\n",
+    "        y_emb = generator.shared_emb(y)\n",
+    "        fake_images = generator(z, y_emb)\n",
+    "        fake_output = discriminator(fake_images.detach(), y)\n",
+    "        \n",
+    "        d_loss = critic_loss(real_output, fake_output)\n",
+    "        d_losses += [d_loss.item()]\n",
+    "        d_loss.backward()\n",
+    "        d_optimizer.step()\n",
+    "        \n",
+    "        for p in discriminator.parameters():\n",
+    "            p.data.clamp_(-0.01, 0.01)\n",
+    "        \n",
+    "        if i % n_critic == 0:\n",
+    "            g_optimizer.zero_grad()\n",
+    "            \n",
+    "            fake_output = discriminator(fake_images, y)\n",
+    "            g_loss = generator_loss(fake_output)\n",
+    "            g_losses += [g_loss.item()]\n",
+    "            \n",
+    "            g_loss.backward()\n",
+    "            g_optimizer.step()\n",
+    "        \n",
+    "        if i % 100 == 0:\n",
+    "            print(f\"Epoch [{epoch}/{num_epochs}], Step [{i}/{len(train_dataloader)}], \"\n",
+    "                  f\"D Loss: {d_loss.item():.4f}, G Loss: {g_loss.item():.4f}\") # type: ignore\n",
+    "        \n",
+    "        # Display generated images every 500 steps\n",
+    "        if i % display_step == 0:\n",
+    "            with torch.no_grad():\n",
+    "                fake_images_display = generator(z, y_emb).detach().cpu()\n",
+    "                real_images_display = real_images.cpu()\n",
+    "                \n",
+    "                # Normalize images to [0, 1] for display\n",
+    "                fake_images_display = (fake_images_display + 1) / 2\n",
+    "                real_images_display = (real_images_display + 1) / 2\n",
+    "                \n",
+    "                # Plot real and fake images side by side\n",
+    "                fig, ax = plt.subplots(1, 2, figsize=(12, 6))\n",
+    "                \n",
+    "                ax[0].imshow(vutils.make_grid(real_images_display, normalize=False).permute(1, 2, 0))\n",
+    "                ax[0].set_title('Real Images')\n",
+    "                ax[0].axis('off')\n",
+    "                \n",
+    "                ax[1].imshow(vutils.make_grid(fake_images_display, normalize=False).permute(1, 2, 0))\n",
+    "                ax[1].set_title('Generated Images')\n",
+    "                ax[1].axis('off')\n",
+    "\n",
+    "                plt.figure(figsize=(10, 5))\n",
+    "                plt.plot(g_losses, label=\"Generator Loss\")\n",
+    "                plt.plot(d_losses, label=\"Discriminator Loss\")\n",
+    "                plt.xlabel(\"Training Steps\")\n",
+    "                plt.ylabel(\"Loss\")\n",
+    "                plt.title(\"Generator and Discriminator Loss During Training\")\n",
+    "                plt.legend()\n",
+    "                plt.grid(True)\n",
+    "                plt.show()\n",
+    "                \n",
+    "                plt.show()\n",
+    "\n",
+    "    # Save checkpoints periodically\n",
+    "    if epoch % 10 == 0:\n",
+    "        torch.save(generator, f'generator_epoch_{epoch}.pth')\n",
+    "        torch.save(discriminator, f'discriminator_epoch_{epoch}.pth')\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}