{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "07d92102",
   "metadata": {},
   "source": [
    "# Lesson 33: Convolutional Neural Networks\n",
    "\n",
    "The layers in Lessons 30-32 treat every input as a flat vector: each pixel gets its own independent weight, with no notion that pixel $(5,5)$ is *near* pixel $(5,6)$. That throws away the entire spatial structure images have, and it means a fully-connected layer must independently re-learn what an edge looks like at every single position. **Convolution** — already built from scratch in Lesson 10 — fixes both problems at once."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da94ca02",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f234b6f",
   "metadata": {},
   "source": [
    "## A conv layer is a projection, restricted and shared\n",
    "\n",
    "Recall Lesson 29's central idea: $y = w^\\top x + b$, a projection. A convolutional layer computes *exactly this*, but for $x$, it only ever looks at a small local patch (say $3\\times3$) instead of the whole image — and, crucially, it reuses the *same* $w$ at every patch location. Two consequences:\n",
    "\n",
    "- **Vastly fewer parameters.** A $3\\times3$ filter has 9 weights, no matter how big the image is — a fully-connected layer over a $16\\times16$ image needs 256 weights *per output unit*.\n",
    "- **Translation equivariance.** Because the same filter slides everywhere, a feature learned at one position is automatically detected at every other position too, for free."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "652f109c",
   "metadata": {},
   "source": [
    "## Conv2d, forward and backward, from scratch\n",
    "\n",
    "This is the same sliding-window computation as Lesson 10's convolution, just with the kernel now a set of *learnable* weights instead of a fixed, hand-designed one — which means it needs a backward pass too."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c02f19dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def conv2d_forward(img, kernel, bias):\n",
    "    N, C, H, W = img.shape\n",
    "    OC, IC, KH, KW = kernel.shape\n",
    "    OH, OW = H - KH + 1, W - KW + 1\n",
    "    out = np.zeros((N, OC, OH, OW))\n",
    "    for n in range(N):\n",
    "        for oc in range(OC):\n",
    "            for i in range(OH):\n",
    "                for j in range(OW):\n",
    "                    patch = img[n, :, i:i + KH, j:j + KW]\n",
    "                    out[n, oc, i, j] = np.sum(patch * kernel[oc]) + bias[oc]\n",
    "    return out\n",
    "\n",
    "def conv2d_backward(img, kernel, grad_out):\n",
    "    N, C, H, W = img.shape\n",
    "    OC, IC, KH, KW = kernel.shape\n",
    "    OH, OW = grad_out.shape[2], grad_out.shape[3]\n",
    "    grad_kernel = np.zeros_like(kernel)\n",
    "    grad_img = np.zeros_like(img)\n",
    "    grad_bias = grad_out.sum(axis=(0, 2, 3))\n",
    "    for n in range(N):\n",
    "        for oc in range(OC):\n",
    "            for i in range(OH):\n",
    "                for j in range(OW):\n",
    "                    patch = img[n, :, i:i + KH, j:j + KW]\n",
    "                    grad_kernel[oc] += grad_out[n, oc, i, j] * patch\n",
    "                    grad_img[n, :, i:i + KH, j:j + KW] += grad_out[n, oc, i, j] * kernel[oc]\n",
    "    return grad_img, grad_kernel, grad_bias"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d402563",
   "metadata": {},
   "source": [
    "### Sanity check against PyTorch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2299f8ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "rng = np.random.default_rng(0)\n",
    "img = rng.normal(size=(1, 1, 8, 8))\n",
    "kernel = rng.normal(size=(1, 1, 3, 3)) * 0.5\n",
    "bias = np.array([0.1])\n",
    "\n",
    "out = conv2d_forward(img, kernel, bias)\n",
    "\n",
    "img_t = torch.tensor(img, requires_grad=True)\n",
    "kernel_t = torch.tensor(kernel, requires_grad=True)\n",
    "bias_t = torch.tensor(bias, requires_grad=True)\n",
    "out_t = F.conv2d(img_t, kernel_t, bias_t)\n",
    "print(f'forward max diff: {np.abs(out - out_t.detach().numpy()).max():.2e}')\n",
    "\n",
    "grad_out = rng.normal(size=out.shape)\n",
    "grad_img, grad_kernel, grad_bias = conv2d_backward(img, kernel, grad_out)\n",
    "out_t.backward(torch.tensor(grad_out))\n",
    "\n",
    "print(f'grad_img max diff:    {np.abs(grad_img - img_t.grad.numpy()).max():.2e}')\n",
    "print(f'grad_kernel max diff: {np.abs(grad_kernel - kernel_t.grad.numpy()).max():.2e}')\n",
    "print(f'grad_bias max diff:   {np.abs(grad_bias - bias_t.grad.numpy()).max():.2e}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2dac3a2c",
   "metadata": {},
   "source": [
    "## Filters you already know: Lesson 12's edge detectors\n",
    "\n",
    "The output of a conv layer before training is nonsense; the point is that gradient descent will *find* useful filters. To see what a useful filter's output already looks like, apply a filter you designed by hand back in Lesson 12: Sobel's edge kernel. Every trained CNN's first-layer filters typically converge to something visually similar to this — oriented edge and blob detectors — independent of what the network was trained to do."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3dd2a2f",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_img = np.zeros((1, 1, 40, 40))\n",
    "test_img[0, 0, 10:30, 10:30] = 1.0\n",
    "\n",
    "sobel_x = np.array([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=np.float64).reshape(1, 1, 3, 3)\n",
    "feature_map = conv2d_forward(test_img, sobel_x, np.zeros(1))\n",
    "\n",
    "fig, axes = plt.subplots(1, 2, figsize=(6, 3.5))\n",
    "axes[0].imshow(test_img[0, 0], cmap='gray')\n",
    "axes[0].set_title('Input')\n",
    "axes[1].imshow(feature_map[0, 0], cmap='gray')\n",
    "axes[1].set_title('Feature map\\n(Sobel filter, as a conv layer)')\n",
    "for ax in axes:\n",
    "    ax.axis('off')\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50195df4",
   "metadata": {},
   "source": [
    "## Pooling: downsampling with a purpose\n",
    "\n",
    "**Max pooling** slides a window over the feature map and keeps only the maximum value in each window, shrinking the spatial size (like the pyramids of Lesson 11) while adding a small amount of local translation invariance — a feature detected at slightly different positions within one pooling window still produces the same output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "086e788f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def maxpool2d(x, size=2):\n",
    "    N, C, H, W = x.shape\n",
    "    OH, OW = H // size, W // size\n",
    "    out = np.zeros((N, C, OH, OW))\n",
    "    for i in range(OH):\n",
    "        for j in range(OW):\n",
    "            patch = x[:, :, i * size:i * size + size, j * size:j * size + size]\n",
    "            out[:, :, i, j] = patch.max(axis=(2, 3))\n",
    "    return out\n",
    "\n",
    "pool_test = rng.normal(size=(1, 1, 8, 8))\n",
    "mine = maxpool2d(pool_test)\n",
    "torch_result = F.max_pool2d(torch.tensor(pool_test), 2)\n",
    "print(f'max pooling max diff vs torch: {np.abs(mine - torch_result.numpy()).max():.2e}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ea0860f",
   "metadata": {},
   "source": [
    "## Why translation equivariance actually matters: a generalization test\n",
    "\n",
    "Build a small classification task: a $16\\times16$ image contains either a plus sign or a circle, at a *random position*. Train on shapes placed only near the center; test on shapes placed only near the corners — positions the model never saw during training. A network that has genuinely learned \"what a plus looks like,\" rather than \"which pixels tend to be on for a plus at these particular training positions,\" should have no trouble with this."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "404ffa93",
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_image(shape_type, cx, cy, size=16):\n",
    "    img = np.zeros((size, size), dtype=np.float32)\n",
    "    if shape_type == 'plus':\n",
    "        img[cy - 1:cy + 2, cx - 3:cx + 4] = 1.0\n",
    "        img[cy - 3:cy + 4, cx - 1:cx + 2] = 1.0\n",
    "    else:\n",
    "        yy, xx = np.mgrid[0:size, 0:size]\n",
    "        img[((xx - cx) ** 2 + (yy - cy) ** 2) <= 9] = 1.0\n",
    "    return img\n",
    "\n",
    "def make_dataset(rng_local, n, position_range):\n",
    "    imgs, labels = [], []\n",
    "    for _ in range(n):\n",
    "        shape_type = rng_local.choice(['plus', 'circle'])\n",
    "        cx, cy = rng_local.integers(*position_range), rng_local.integers(*position_range)\n",
    "        imgs.append(make_image(shape_type, cx, cy))\n",
    "        labels.append(0.0 if shape_type == 'plus' else 1.0)\n",
    "    return np.array(imgs, dtype=np.float32), np.array(labels, dtype=np.float32)\n",
    "\n",
    "data_rng = np.random.default_rng(1)\n",
    "X_train, y_train = make_dataset(data_rng, 300, (5, 11))   # near center\n",
    "X_test, y_test = make_dataset(data_rng, 150, (3, 5))       # near a corner -- never seen in training\n",
    "\n",
    "fig, axes = plt.subplots(1, 4, figsize=(9, 2.5))\n",
    "for ax, im, title in zip(axes, [X_train[0], X_train[1], X_test[0], X_test[1]],\n",
    "                          ['train example', 'train example', 'test example\\n(unseen position)', 'test example\\n(unseen position)']):\n",
    "    ax.imshow(im, cmap='gray')\n",
    "    ax.set_title(title, fontsize=8)\n",
    "    ax.axis('off')\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68166338",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MLPClassifier(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.net = nn.Sequential(nn.Flatten(), nn.Linear(16 * 16, 32), nn.ReLU(), nn.Linear(32, 1))\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.net(x).squeeze(-1)\n",
    "\n",
    "class CNNClassifier(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.conv = nn.Sequential(\n",
    "            nn.Conv2d(1, 8, 5, padding=2), nn.ReLU(),\n",
    "            nn.MaxPool2d(2),\n",
    "            nn.Conv2d(8, 16, 5, padding=2), nn.ReLU(),\n",
    "            nn.AdaptiveMaxPool2d(1),  # global max pool: collapses ALL spatial position info\n",
    "        )\n",
    "        self.fc = nn.Linear(16, 1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.fc(self.conv(x).flatten(1)).squeeze(-1)\n",
    "\n",
    "def train_and_eval(model_cls, Xtr, ytr, Xte, yte, seed, epochs=400, lr=0.01):\n",
    "    torch.manual_seed(seed)  # seed BEFORE constructing the model, so init is actually reproducible\n",
    "    model = model_cls()\n",
    "    opt = torch.optim.Adam(model.parameters(), lr=lr)\n",
    "    for _ in range(epochs):\n",
    "        opt.zero_grad()\n",
    "        loss = F.binary_cross_entropy_with_logits(model(Xtr), ytr)\n",
    "        loss.backward()\n",
    "        opt.step()\n",
    "    with torch.no_grad():\n",
    "        train_acc = ((model(Xtr) > 0).float() == ytr).float().mean().item()\n",
    "        test_acc = ((model(Xte) > 0).float() == yte).float().mean().item()\n",
    "    return train_acc, test_acc\n",
    "\n",
    "Xtr_t = torch.tensor(X_train).unsqueeze(1)\n",
    "ytr_t = torch.tensor(y_train)\n",
    "Xte_t = torch.tensor(X_test).unsqueeze(1)\n",
    "yte_t = torch.tensor(y_test)\n",
    "\n",
    "mlp_train_acc, mlp_test_acc = train_and_eval(MLPClassifier, Xtr_t, ytr_t, Xte_t, yte_t, seed=0)\n",
    "cnn_train_acc, cnn_test_acc = train_and_eval(CNNClassifier, Xtr_t, ytr_t, Xte_t, yte_t, seed=0)\n",
    "\n",
    "print(f'{\"model\":>6} {\"train acc\":>10} {\"test acc (unseen positions)\":>30}')\n",
    "print(f'{\"MLP\":>6} {mlp_train_acc:>10.1%} {mlp_test_acc:>30.1%}')\n",
    "print(f'{\"CNN\":>6} {cnn_train_acc:>10.1%} {cnn_test_acc:>30.1%}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c4fbedc",
   "metadata": {},
   "source": [
    "Both models fit the training data perfectly. On positions neither has ever seen, the flatten-based MLP does little better than a coin flip &mdash; it memorized *which pixels* tend to be on for each class at the training positions, and that knowledge doesn't transfer. The CNN, whose global max pooling forces the final decision to depend only on *what filters fired somewhere*, not *where*, generalizes to the new positions with no loss in accuracy at all."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "adea55e0",
   "metadata": {},
   "source": [
    "### Exercise\n",
    "\n",
    "1. Replace `nn.AdaptiveMaxPool2d(1)` in `CNNClassifier` with `nn.Flatten()` directly on the conv features (removing the global pooling, so spatial position is preserved all the way to the final linear layer). Retrain and re-evaluate on the unseen-position test set. Does the CNN's generalization advantage survive?\n",
    "2. Increase the gap between train and test position ranges (e.g. train on `(5, 11)`, test on `(0, 3)`, right at the image border where shapes get clipped). Does the CNN's accuracy hold up, or does it degrade &mdash; and if it degrades, is that a translation-invariance failure or something else entirely (think about what happens to a shape's *appearance*, not just its position, right at an edge)?\n",
    "3. `conv2d_forward` above pads nothing (`OH = H - KH + 1`), so the output shrinks with every layer. Modify it to support zero-padding (pad the input by `KH//2` on each side before sliding, matching PyTorch's `padding=` argument) and confirm the output size matches the input size, validated against `F.conv2d(..., padding=1)`."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.x"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}