{
 "nbformat": 4,
 "nbformat_minor": 5,
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.11"
  },
  "colab": {
   "provenance": []
  }
 },
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cell-0000",
   "metadata": {},
   "source": [
    "# Generative Adversarial Networks\n",
    "\n",
    "Two networks locked in a forgery contest — train a GAN from a 1-D toy demo in your browser up to a DCGAN that draws handwritten digits.\n",
    "\n",
    "*Part of the free [Deep Learning with PyTorch](https://ramadnsyh.dev/courses/deep-learning) course by [Muhammad Ramadiansyah](https://ramadnsyh.dev). This notebook is generated from the interactive lesson — [read it online](https://ramadnsyh.dev/courses/deep-learning/gans).*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0001",
   "metadata": {},
   "source": [
    "The autoencoder learned to compress and redraw data it had seen. A GAN goes\n",
    "further: it learns to draw data that **never existed** — faces, digits,\n",
    "artwork — by turning generation into a game between two networks. In this\n",
    "final lesson you'll build the adversarial intuition with a demo small enough\n",
    "to run in your browser, then scale the same idea up to a DCGAN that generates\n",
    "handwritten digits."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0002",
   "metadata": {},
   "source": [
    "## The counterfeiter and the detective\n",
    "\n",
    "A GAN (generative adversarial network, Goodfellow et al., 2014) trains two\n",
    "networks against each other:\n",
    "\n",
    "- The **generator G** is a counterfeiter. It takes a random noise vector `z`\n",
    "  and transforms it into a fake sample, `G(z)` — a fake image, say. It never\n",
    "  sees real data directly.\n",
    "- The **discriminator D** is a detective. Given a sample, it outputs the\n",
    "  probability that the sample is real. It trains on both real data (label 1)\n",
    "  and the generator's fakes (label 0).\n",
    "\n",
    "They improve *because of each other*. Early on, the fakes are garbage and the\n",
    "detective wins easily. But the detective's verdicts flow back through\n",
    "backpropagation as a training signal for the counterfeiter: \"this fake was\n",
    "spotted because of these pixels.\" The counterfeiter adjusts, the fakes get\n",
    "better, the detective is forced to sharpen its criteria, and around it goes.\n",
    "At the theoretical equilibrium, the fakes are indistinguishable from real\n",
    "data and the detective is reduced to guessing — 50/50.\n",
    "\n",
    "Formally the game is a **minimax objective**:\n",
    "`min_G max_D  E[log D(x)] + E[log(1 − D(G(z)))]`. In plain words: the\n",
    "discriminator tunes its weights to *maximize* its accuracy — assign high\n",
    "`D(x)` to real samples and low `D(G(z))` to fakes — while the generator tunes\n",
    "its weights to *minimize* that same score by making fakes the discriminator\n",
    "scores high. One value function, two players pulling in opposite directions."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0003",
   "metadata": {},
   "source": [
    "> **The non-saturating trick**\n",
    "> \n",
    "> In practice the generator doesn't minimize `log(1 − D(G(z)))` — that gradient\n",
    "> vanishes exactly when the generator is losing badly. Instead it *maximizes*\n",
    "> `log D(G(z))`, which gives strong gradients when fakes are easily spotted.\n",
    "> Same game, healthier learning signal — and it's what every implementation,\n",
    "> including ours below, actually uses."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0004",
   "metadata": {},
   "source": [
    "## The adversarial game in 1-D, live\n",
    "\n",
    "Images are too big to watch a GAN think, so let's shrink the problem until\n",
    "every moving part is visible. The \"real data\" is just numbers drawn from\n",
    "`N(4, 0.5)`. The generator is the simplest possible one — it reshapes standard\n",
    "Gaussian noise as `x = mu + sigma * z`, so its only weights are `mu` and\n",
    "`sigma`. The discriminator is a tiny logistic regression on the features\n",
    "`x` and `x²` (quadratic features can perfectly separate two Gaussians). Both\n",
    "are trained with plain gradient steps, alternating — a real GAN loop, in\n",
    "numpy:"
   ]
  },
  {
   "cell_type": "code",
   "id": "cell-0005",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "rng = np.random.default_rng(0)\n",
    "\n",
    "def sigmoid(t):\n",
    "    return 1 / (1 + np.exp(-np.clip(t, -30, 30)))\n",
    "\n",
    "def feats(x):   # discriminator features: [1, x, x^2]\n",
    "    return np.stack([np.ones_like(x), x, x**2], axis=1)\n",
    "\n",
    "w = np.zeros(3)          # discriminator weights\n",
    "mu, sigma = 0.0, 1.0     # generator: fake = mu + sigma * z\n",
    "lr_d, lr_g, n = 0.05, 0.05, 128\n",
    "\n",
    "for step in range(1, 601):\n",
    "    real = rng.normal(4.0, 0.5, n)\n",
    "    z = rng.normal(0.0, 1.0, n)\n",
    "    fake = mu + sigma * z\n",
    "\n",
    "    # --- discriminator step: real -> 1, fake -> 0 ---\n",
    "    for xb, label in [(real, 1.0), (fake, 0.0)]:\n",
    "        p = sigmoid(feats(xb) @ w)\n",
    "        w -= lr_d * feats(xb).T @ (p - label) / n\n",
    "\n",
    "    # --- generator step: move fakes where D says 'real' ---\n",
    "    p_fake = sigmoid(feats(fake) @ w)\n",
    "    dlogit = w[1] + 2 * w[2] * fake        # how D's logit changes with x\n",
    "    upstream = (1 - p_fake) * dlogit       # gradient of log D(fake)\n",
    "    mu += lr_g * np.mean(upstream)         # gradient ASCENT\n",
    "    sigma = max(0.05, sigma + lr_g * np.mean(upstream * z))\n",
    "\n",
    "    if step % 100 == 0:\n",
    "        print(f\"step {step:3d}: mu={mu:.2f}  sigma={sigma:.2f}  \"\n",
    "              f\"D(fake)={np.mean(p_fake):.2f}\")\n",
    "\n",
    "real = rng.normal(4.0, 0.5, 3000)\n",
    "fake = mu + sigma * rng.normal(0.0, 1.0, 3000)\n",
    "plt.hist(real, bins=60, alpha=0.6, label=\"real  N(4, 0.5)\")\n",
    "plt.hist(fake, bins=60, alpha=0.6, label=f\"generated  N({mu:.1f}, {sigma:.1f})\")\n",
    "plt.legend(); plt.title(\"Generated distribution after adversarial training\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0006",
   "metadata": {},
   "source": [
    "The generator starts producing numbers around 0 and, guided only by the\n",
    "discriminator's verdicts, marches its distribution over to sit on top of the\n",
    "real one. Look honestly at the final numbers, though: `mu` lands near 4 but\n",
    "`sigma` settles around 0.7 rather than 0.5, and `D(fake)` oscillates instead\n",
    "of resting at 0.5. Even in one dimension with four parameters, the game\n",
    "circles the equilibrium rather than settling on it — remember this when we\n",
    "discuss instability."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0007",
   "metadata": {},
   "source": [
    "## Scaling up: DCGAN on MNIST\n",
    "\n",
    "For images, both players become convolutional networks — the **DCGAN**\n",
    "recipe. The generator runs a CNN in reverse: `ConvTranspose2d` layers\n",
    "*upsample* a noise vector into an image, doubling resolution at each step.\n",
    "The discriminator is an ordinary CNN classifier ending in a single logit.\n",
    "This needs a GPU — run it in the notebook on Colab."
   ]
  },
  {
   "cell_type": "code",
   "id": "cell-0008",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision import datasets, transforms\n",
    "\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "\n",
    "# scale images to [-1, 1] to match the generator's Tanh output\n",
    "tfm = transforms.Compose([transforms.ToTensor(),\n",
    "                          transforms.Normalize(0.5, 0.5)])\n",
    "train_set = datasets.MNIST(\"data\", train=True, download=True, transform=tfm)\n",
    "loader = DataLoader(train_set, batch_size=128, shuffle=True)\n",
    "\n",
    "z_dim = 100\n",
    "\n",
    "generator = nn.Sequential(\n",
    "    # (batch, 100, 1, 1) -> (batch, 128, 7, 7)\n",
    "    nn.ConvTranspose2d(z_dim, 128, kernel_size=7, stride=1, padding=0),\n",
    "    nn.BatchNorm2d(128), nn.ReLU(),\n",
    "    # -> (batch, 64, 14, 14)\n",
    "    nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1),\n",
    "    nn.BatchNorm2d(64), nn.ReLU(),\n",
    "    # -> (batch, 1, 28, 28)\n",
    "    nn.ConvTranspose2d(64, 1, kernel_size=4, stride=2, padding=1),\n",
    "    nn.Tanh(),\n",
    ").to(device)\n",
    "\n",
    "discriminator = nn.Sequential(\n",
    "    # (batch, 1, 28, 28) -> (batch, 64, 14, 14)\n",
    "    nn.Conv2d(1, 64, kernel_size=4, stride=2, padding=1),\n",
    "    nn.LeakyReLU(0.2),\n",
    "    # -> (batch, 128, 7, 7)\n",
    "    nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1),\n",
    "    nn.BatchNorm2d(128), nn.LeakyReLU(0.2),\n",
    "    # -> (batch, 1, 1, 1) -> one logit per image\n",
    "    nn.Conv2d(128, 1, kernel_size=7, stride=1, padding=0),\n",
    "    nn.Flatten(),\n",
    ").to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0009",
   "metadata": {},
   "source": [
    "The training loop alternates the two players every batch. Note the\n",
    "`.detach()` when training the discriminator (no generator gradients wanted)\n",
    "and the flipped labels when training the generator:"
   ]
  },
  {
   "cell_type": "code",
   "id": "cell-0010",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "criterion = nn.BCEWithLogitsLoss()\n",
    "opt_d = torch.optim.Adam(discriminator.parameters(), lr=2e-4, betas=(0.5, 0.999))\n",
    "opt_g = torch.optim.Adam(generator.parameters(), lr=2e-4, betas=(0.5, 0.999))\n",
    "\n",
    "fixed_z = torch.randn(64, z_dim, 1, 1, device=device)  # to watch progress\n",
    "\n",
    "for epoch in range(25):\n",
    "    for real, _ in loader:\n",
    "        real = real.to(device)\n",
    "        bs = len(real)\n",
    "        ones, zeros = torch.ones(bs, 1, device=device), torch.zeros(bs, 1, device=device)\n",
    "\n",
    "        # --- 1. discriminator step ---\n",
    "        z = torch.randn(bs, z_dim, 1, 1, device=device)\n",
    "        fake = generator(z)\n",
    "        loss_d = (criterion(discriminator(real), ones) +\n",
    "                  criterion(discriminator(fake.detach()), zeros))\n",
    "        opt_d.zero_grad(); loss_d.backward(); opt_d.step()\n",
    "\n",
    "        # --- 2. generator step: make D say 'real' on fakes ---\n",
    "        loss_g = criterion(discriminator(fake), ones)\n",
    "        opt_g.zero_grad(); loss_g.backward(); opt_g.step()\n",
    "\n",
    "    print(f\"epoch {epoch + 1}: loss_D={loss_d.item():.3f} loss_G={loss_g.item():.3f}\")\n",
    "\n",
    "    if (epoch + 1) % 5 == 0:            # sample a grid every few epochs\n",
    "        from torchvision.utils import make_grid\n",
    "        import matplotlib.pyplot as plt\n",
    "        with torch.no_grad():\n",
    "            samples = generator(fixed_z).cpu() * 0.5 + 0.5   # back to [0, 1]\n",
    "        plt.figure(figsize=(6, 6))\n",
    "        plt.imshow(make_grid(samples, nrow=8).permute(1, 2, 0), cmap=\"gray\")\n",
    "        plt.axis(\"off\"); plt.title(f\"epoch {epoch + 1}\")\n",
    "        plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0011",
   "metadata": {},
   "source": [
    "Because `fixed_z` never changes, the periodic grids show the *same* 64 noise\n",
    "vectors maturing from static, to blobs, to recognizable digits over about 25\n",
    "epochs. Every digit in the final grid is a drawing that exists nowhere in\n",
    "MNIST."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0012",
   "metadata": {},
   "source": [
    "## Why GAN training is famously unstable\n",
    "\n",
    "Unlike every loss you've minimized so far, a GAN has no single number going\n",
    "reliably down — it's two losses chasing each other, and the \"landscape\" moves\n",
    "whenever either player does. The classic failure modes:\n",
    "\n",
    "- **Mode collapse.** The generator finds one output that reliably fools the\n",
    "  discriminator and produces only that — a GAN that draws convincing 1s and\n",
    "  nothing else. Diversity dies because the objective never explicitly demands\n",
    "  it. (Our 1-D demo showed a cousin of this: on two-mode data, a too-simple\n",
    "  generator drifts to one mode or smears across both.)\n",
    "- **Non-convergence.** The players can circle each other forever — generator\n",
    "  adapts, discriminator re-adapts, losses oscillate — without approaching\n",
    "  equilibrium, exactly like the residual wobble in the 1-D demo.\n",
    "- **Imbalance.** A discriminator that wins too hard gives the generator\n",
    "  near-zero gradients; one that's too weak gives it meaningless guidance.\n",
    "\n",
    "The battle-tested tricks, most from the DCGAN paper, and already baked into\n",
    "the code above: learning rate `2e-4` with Adam betas `(0.5, 0.999)` (the\n",
    "lower momentum term damps the oscillation), `BatchNorm` in both networks,\n",
    "`LeakyReLU` in the discriminator, `Tanh` output with inputs normalized to\n",
    "`[-1, 1]`, and strided convolutions instead of pooling. Later research added\n",
    "better objectives (Wasserstein loss, spectral normalization) attacking the\n",
    "same instabilities."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0013",
   "metadata": {},
   "source": [
    "## GANs today\n",
    "\n",
    "Honest modern context: for image generation, **diffusion models** — the\n",
    "engines behind Stable Diffusion and friends — have largely displaced GANs.\n",
    "They optimize a plain denoising objective (an idea you already met in the\n",
    "denoising autoencoder!), which sidesteps the two-player instability and\n",
    "covers modes much more reliably. But GANs remain in real use where speed\n",
    "matters — a GAN generates in one forward pass versus a diffusion model's many\n",
    "denoising steps — for super-resolution, image-to-image translation, and as\n",
    "adversarial *components*: many state-of-the-art systems still bolt on a\n",
    "discriminator as an extra \"does this look real?\" loss. The adversarial idea\n",
    "outlived the architecture."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0014",
   "metadata": {},
   "source": [
    "### 🏋️ Exercise — Watch a generator struggle with two modes\n",
    "\n",
    "In the 1-D PyRunner above, make the real data **bimodal**: half the samples\n",
    "from `N(0, 0.5)` and half from `N(6, 0.5)`. Rerun the training and look at\n",
    "the final histogram and the learned `mu` and `sigma`. Where did the generator\n",
    "put its probability mass, and why is this a miniature version of mode\n",
    "collapse? Write two sentences explaining what the generator would need in\n",
    "order to fix it."
   ]
  },
  {
   "cell_type": "code",
   "id": "cell-0015",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Your solution here"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0016",
   "metadata": {},
   "source": [
    "<details><summary>✅ Show solution</summary>\n",
    "\n",
    "```python\n",
    "# In the 1-D PyRunner, replace both real-data draws with a mixture:\n",
    "#\n",
    "#     modes = rng.choice([0.0, 6.0], n)\n",
    "#     real = rng.normal(modes, 0.5)\n",
    "#\n",
    "# (and the same for the 3000-sample histogram at the end).\n",
    "#\n",
    "# Typical outcome: mu drifts toward one of the modes (around 1.2 in one\n",
    "# run) and sigma inflates to about 2.4 as the generator smears itself\n",
    "# across the gap - placing lots of samples near x = 3 where NO real\n",
    "# data lives. The generator family (one Gaussian) simply cannot\n",
    "# represent two modes, so it must either collapse onto one mode or\n",
    "# blur across both. This is the small-scale cousin of mode collapse:\n",
    "# when the generator cannot or does not cover every mode of the data,\n",
    "# the adversarial game gives no explicit penalty for ignoring some.\n",
    "```\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0017",
   "metadata": {},
   "source": [
    "That's a wrap — not just on generative models, but on the entire deep\n",
    "learning course. You started with a single neuron and ended by training two\n",
    "networks to out-scheme each other into drawing digits from pure noise:\n",
    "backpropagation, CNNs, transfer learning, RNNs, embeddings, autoencoders, and\n",
    "GANs are all in your toolkit now. Congratulations on finishing! When you're\n",
    "ready for more, head back to [the course catalog](/courses) and pick your\n",
    "next adventure."
   ]
  }
 ]
}