dcgan.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "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",
    "import torchvision.transforms as transforms\n",
    "from torch.utils.data import DataLoader\n",
    "from torch.utils.tensorboard import SummaryWriter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Discriminator(nn.Module):\n",
    "    def __init__(self, channels_img, features_d): # channels_img: 3 for rgb images  features_d: 64\n",
    "        super(Discriminator, self).__init__()\n",
    "        self.disc = nn.Sequential(\n",
    "            # input: N x channels_img x 64 x 64\n",
    "            nn.Conv2d(\n",
    "                channels_img, features_d, kernel_size=4, stride=2, padding=1\n",
    "            ),\n",
    "            nn.LeakyReLU(0.2),\n",
    "            # _block(in_channels, out_channels, kernel_size, stride, padding)\n",
    "            self._block(features_d, features_d * 2, 4, 2, 1),\n",
    "            self._block(features_d * 2, features_d * 4, 4, 2, 1),\n",
    "            self._block(features_d * 4, features_d * 8, 4, 2, 1),\n",
    "            # After all _block img output is 4x4 (Conv2d below makes into 1x1)\n",
    "            nn.Conv2d(features_d * 8, 1, kernel_size=4, stride=2, padding=0),\n",
    "            nn.Sigmoid(),\n",
    "        )\n",
    "\n",
    "    def _block(self, in_channels, out_channels, kernel_size, stride, padding):\n",
    "        return nn.Sequential(\n",
    "            nn.Conv2d(\n",
    "                in_channels,\n",
    "                out_channels,\n",
    "                kernel_size,\n",
    "                stride,\n",
    "                padding,\n",
    "                bias=False,\n",
    "            ),\n",
    "            #nn.BatchNorm2d(out_channels),\n",
    "            nn.LeakyReLU(0.2),\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.disc(x)\n",
    "\n",
    "\n",
    "class Generator(nn.Module):\n",
    "    def __init__(self, channels_noise, channels_img, features_g): # channels_noise: 100, channels_img: 3, features_g: 64\n",
    "        super(Generator, self).__init__()\n",
    "        self.net = nn.Sequential(\n",
    "            # Input: N x channels_noise x 1 x 1\n",
    "            self._block(channels_noise, features_g * 16, 4, 1, 0),  # img: 4x4\n",
    "            self._block(features_g * 16, features_g * 8, 4, 2, 1),  # img: 8x8\n",
    "            self._block(features_g * 8, features_g * 4, 4, 2, 1),  # img: 16x16\n",
    "            self._block(features_g * 4, features_g * 2, 4, 2, 1),  # img: 32x32\n",
    "            nn.ConvTranspose2d(\n",
    "                features_g * 2, channels_img, kernel_size=4, stride=2, padding=1\n",
    "            ),\n",
    "            # Output: N x channels_img x 64 x 64\n",
    "            nn.Tanh(),\n",
    "        )\n",
    "\n",
    "    def _block(self, in_channels, out_channels, kernel_size, stride, padding):\n",
    "        return nn.Sequential(\n",
    "            nn.ConvTranspose2d(\n",
    "                in_channels,\n",
    "                out_channels,\n",
    "                kernel_size,\n",
    "                stride,\n",
    "                padding,\n",
    "                bias=False,\n",
    "            ),\n",
    "            #nn.BatchNorm2d(out_channels),\n",
    "            nn.ReLU(),\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.net(x)\n",
    "\n",
    "\n",
    "def initialize_weights(model):\n",
    "    # Initializes weights according to the DCGAN paper\n",
    "    for m in model.modules():\n",
    "        if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d, nn.BatchNorm2d)):\n",
    "            nn.init.normal_(m.weight.data, 0.0, 0.02)\n",
    "\n",
    "def test():\n",
    "    N, in_channels, H, W = 8, 3, 64, 64\n",
    "    noise_dim = 100\n",
    "    x = torch.randn((N, in_channels, H, W))\n",
    "    disc = Discriminator(in_channels, 8)\n",
    "    assert disc(x).shape == (N, 1, 1, 1), \"Discriminator test failed\"\n",
    "    gen = Generator(noise_dim, in_channels, 8)\n",
    "    z = torch.randn((N, noise_dim, 1, 1))\n",
    "    assert gen(z).shape == (N, in_channels, H, W), \"Generator test failed\"\n",
    "\n",
    "\n",
    "# test()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Hyperparameters etc.\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "LEARNING_RATE = 2e-4  # could also use two lrs, one for gen and one for disc\n",
    "BATCH_SIZE = 128\n",
    "IMAGE_SIZE = 64\n",
    "CHANNELS_IMG = 1 # If you train on MNIST, remember to set channels_img to 1\n",
    "NOISE_DIM = 100\n",
    "NUM_EPOCHS = 5\n",
    "FEATURES_DISC = 64\n",
    "FEATURES_GEN = 64"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "stats = (0.5, 0.5, 0.5), (0.5, 0.5, 0.5)\n",
    "transforms = transforms.Compose(\n",
    "    [\n",
    "        transforms.Resize(IMAGE_SIZE),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize(*stats),\n",
    "    ]\n",
    ")\n",
    "\n",
    "def denorm(img_tensors):\n",
    "    return img_tensors * stats[1][0] + stats[0][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = datasets.MNIST(root=\"dataset/\", train=True, transform=transforms,\n",
    "                       download=True)\n",
    "\n",
    "#dataset = datasets.ImageFolder(root=\"celeb_dataset\", transform=transforms)\n",
    "dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)\n",
    "gen = Generator(NOISE_DIM, CHANNELS_IMG, FEATURES_GEN).to(device)\n",
    "disc = Discriminator(CHANNELS_IMG, FEATURES_DISC).to(device)\n",
    "initialize_weights(gen)\n",
    "initialize_weights(disc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "opt_gen = optim.Adam(gen.parameters(), lr=LEARNING_RATE, betas=(0.5, 0.999))\n",
    "opt_disc = optim.Adam(disc.parameters(), lr=LEARNING_RATE, betas=(0.5, 0.999))\n",
    "criterion = nn.BCELoss()\n",
    "\n",
    "fixed_noise = torch.randn(32, NOISE_DIM, 1, 1).to(device)\n",
    "writer_real = SummaryWriter(f\"logs/real\")\n",
    "writer_fake = SummaryWriter(f\"logs/fake\")\n",
    "step = 0\n",
    "\n",
    "gen.train()\n",
    "disc.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for epoch in range(NUM_EPOCHS):\n",
    "    # Target labels not needed! <3 unsupervised\n",
    "    for batch_idx, (real, _) in enumerate(dataloader):\n",
    "        real = real.to(device)\n",
    "        noise = torch.randn(BATCH_SIZE, NOISE_DIM, 1, 1).to(device)\n",
    "        fake = gen(noise)\n",
    "\n",
    "        ### Train Discriminator: max log(D(x)) + log(1 - D(G(z)))\n",
    "        disc_real = disc(real).reshape(-1)\n",
    "        loss_disc_real = criterion(disc_real, torch.ones_like(disc_real))\n",
    "        disc_fake = disc(fake.detach()).reshape(-1)\n",
    "        loss_disc_fake = criterion(disc_fake, torch.zeros_like(disc_fake))\n",
    "        loss_disc = (loss_disc_real + loss_disc_fake) / 2\n",
    "        disc.zero_grad()\n",
    "        loss_disc.backward()\n",
    "        opt_disc.step()\n",
    "\n",
    "        # Train Generator: min log(1 - D(G(z))) <-> max log(D(G(z))\n",
    "        output = disc(fake).reshape(-1)\n",
    "        loss_gen = criterion(output, torch.ones_like(output))\n",
    "        gen.zero_grad()\n",
    "        loss_gen.backward()\n",
    "        opt_gen.step()\n",
    "\n",
    "        # Print losses occasionally and print to tensorboard\n",
    "        if batch_idx % 100 == 0:\n",
    "            print(\n",
    "                f\"Epoch [{epoch}/{NUM_EPOCHS}] Batch {batch_idx}/{len(dataloader)} \\\n",
    "                  Loss D: {loss_disc:.4f}, loss G: {loss_gen:.4f}\"\n",
    "            )\n",
    "\n",
    "            with torch.no_grad():\n",
    "                fake = gen(fixed_noise)\n",
    "                # take out (up to) 32 examples\n",
    "                img_grid_real = torchvision.utils.make_grid(\n",
    "                    real[:32], normalize=True\n",
    "                )\n",
    "                img_grid_fake = torchvision.utils.make_grid(\n",
    "                    fake[:32], normalize=True\n",
    "                )\n",
    "\n",
    "                writer_real.add_image(\"Real\", img_grid_real, global_step=step)\n",
    "                writer_fake.add_image(\"Fake\", img_grid_fake, global_step=step)\n",
    "\n",
    "            step += 1\n",
    "\n"
   ]
  }
 ],
 "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.9.13 (tags/v3.9.13:6de2ca5, May 17 2022, 16:36:42) [MSC v.1929 64 bit (AMD64)]"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "3226585472a06c598dfa09e82eb7bd2822a77ef3b52edcc2f550a8000bbf5d90"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}