SimpleGAN.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torchvision\n",
    "import torchvision.datasets as datasets\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "import torchvision.transforms as transforms\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "class Discriminator(nn.Module):\n",
    "    def __init__(self,img_dim):\n",
    "        super().__init__()\n",
    "        self.disc = nn.Sequential(\n",
    "            nn.Linear(img_dim,128),\n",
    "            nn.LeakyReLU(0.1),\n",
    "            nn.Linear(128,1),\n",
    "            nn.Sigmoid()\n",
    "             )\n",
    "    def forward(self,x):\n",
    "        return self.disc(x)\n",
    "\n",
    "class Generator(nn.Module):\n",
    "    def __init__(self,z_dim,img_dim):\n",
    "        super().__init__()\n",
    "        self.gen = nn.Sequential(\n",
    "            nn.Linear(z_dim,256),\n",
    "            nn.LeakyReLU(0.1),\n",
    "            nn.Linear(256,img_dim),\n",
    "            nn.Tanh()\n",
    "              )\n",
    "\n",
    "    def forward(self,x):\n",
    "         return self.gen(x)\n",
    "\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "lr = 5e-4\n",
    "z_dim = 64\n",
    "image_dim = 28 * 28 * 1\n",
    "batch_size = 32\n",
    "num_epochs = 5\n",
    "\n",
    "disc = Discriminator(image_dim).to(device)\n",
    "gen = Generator(z_dim,image_dim).to(device)\n",
    "fixed_noise = torch.randn(batch_size,z_dim).to(device)\n",
    "transform = transforms.Compose(\n",
    "[transforms.ToTensor(),transforms.Normalize((0.5,),(0.5,))]\n",
    ")\n",
    "dataset = datasets.MNIST(root = \"dataset/\", transform = transform,download=True)\n",
    "loader = DataLoader(dataset, batch_size=batch_size,shuffle=True)\n",
    "opt_disc = optim.Adam(disc.parameters(),lr =lr)\n",
    "opt_gen = optim.Adam(gen.parameters(),lr = lr)\n",
    "criterion = nn.BCELoss()\n",
    "writer_fake = SummaryWriter(f\"runs/GAN_MNIST/fake\")\n",
    "writer_real = SummaryWriter(f\"runs/GAN_MNIST/real\")\n",
    "step = 0\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    for batch_idx, (real,_) in enumerate(loader):\n",
    "            real = real.view(-1,784).to(device)\n",
    "            batch_size = real.shape[0]\n",
    "\n",
    "            noise = torch.randn((batch_size,z_dim)).to(device)\n",
    "            fake = gen(noise)\n",
    "            disc_real = disc(real).view(-1)\n",
    "            lossD_real = criterion(disc_real,torch.ones_like(disc_real))\n",
    "            disc_fake = disc(fake).view(-1)\n",
    "            lossD_fake = criterion(disc_fake,torch.zeros_like(disc_fake))\n",
    "            lossD = (lossD_fake + lossD_real)/2\n",
    "            disc.zero_grad()\n",
    "            lossD.backward(retain_graph = True)\n",
    "\n",
    "            opt_disc.step()\n",
    "\n",
    "            output = disc(fake).view(-1)\n",
    "            lossG = criterion(output,torch.ones_like(output))\n",
    "            gen.zero_grad()\n",
    "            lossG.backward()\n",
    "            opt_gen.step()\n",
    "\n",
    "            if batch_idx == 0:\n",
    "                        print(\n",
    "                            f\"Epoch [{epoch}/{num_epochs}] Batch {batch_idx}/{len(loader)} \\\n",
    "                                  Loss D: {lossD:.4f}, loss G: {lossG:.4f}\"\n",
    "                        )\n",
    "\n",
    "                        with torch.no_grad():\n",
    "                            fake = gen(fixed_noise).reshape(-1, 1, 28, 28)\n",
    "                            data = real.reshape(-1, 1, 28, 28)\n",
    "\n",
    "                            img_grid_fake = torchvision.utils.make_grid(fake, normalize=True)\n",
    "                            img_grid_real = torchvision.utils.make_grid(data, normalize=True)\n",
    "\n",
    "                            # print(type(img_grid_fake.cpu().detach().numpy()))\n",
    "                            # writer_fake.add_image(\n",
    "                            #     \"Mnist Fake Images\", img_grid_fake, global_step=step\n",
    "                            # )\n",
    "                            # writer_real.add_image(\n",
    "                            #     \"Mnist Real Images\", img_grid_real, global_step=step\n",
    "                            # )\n",
    "                            plt.imshow(img_grid_fake.cpu().permute(1, 2, 0))\n",
    "\n",
    "\n",
    "                            step += 1"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}