{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b67ebcc8",
   "metadata": {},
   "source": [
    "# Convolutional networks: from one filter to ResNets, detection and transfer learning\n",
    "\n",
    "Convolution by hand and in PyTorch, output sizes, pooling, a CNN against a plain network, a residual block, depthwise separable convolutions, augmentation with torchvision.transforms.v2, IoU and non-max suppression, transfer learning with torchvision's pre-trained weights, triplet loss and the style (Gram) matrix.\n",
    "\n",
    "[Open in Google Colab](https://colab.research.google.com/github/mahmoudyosrimahmoud13/inpace-library/blob/main/courses/dl-4-cnns/files/dl4_convolutional_networks_pytorch.ipynb) · part of the inPace course library. Slides © DeepLearning.AI (CC BY-SA 2.0); this notebook: inPace (CC BY-SA 4.0).\n",
    "\n",
    "**Run it top to bottom** (Runtime → Run all in Colab). Each part says what to look for."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0f738d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Works in Google Colab as it is (PyTorch is already installed there).\n",
    "# On your own computer:  pip install torch torchvision scikit-learn scipy matplotlib pandas\n",
    "# One framework per notebook: PyTorch and TensorFlow can clash when loaded into the same Python process.\n",
    "import numpy as np, matplotlib.pyplot as plt, torch\n",
    "from torch import nn\n",
    "torch.manual_seed(0); np.random.seed(0)\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"mps\" if torch.backends.mps.is_available() else \"cpu\"\n",
    "print(\"PyTorch\", torch.__version__, \"| device:\", device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c31a1918",
   "metadata": {},
   "source": [
    "## 1. Convolution by hand, then with PyTorch\n",
    "A vertical-edge filter on a 6×6 image (the slide's example). PyTorch's `conv2d` computes the same thing (deep learning's 'convolution' is cross-correlation: no flipping)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09b22555",
   "metadata": {},
   "outputs": [],
   "source": [
    "img = np.array([[10, 10, 10, 0, 0, 0]] * 6, dtype=np.float32)\n",
    "k = np.array([[1, 0, -1]] * 3, dtype=np.float32)\n",
    "by_hand = np.array([[(img[i:i+3, j:j+3] * k).sum() for j in range(4)] for i in range(4)])\n",
    "torch_out = nn.functional.conv2d(torch.tensor(img)[None, None], torch.tensor(k)[None, None])[0, 0].numpy()\n",
    "print(by_hand); print(\"same as PyTorch:\", np.allclose(by_hand, torch_out))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "30dcc333",
   "metadata": {},
   "source": [
    "## 2. Output size: ⌊(n + 2p − f) / s⌋ + 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd1440a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def out_size(n, f, p=0, s=1): return (n + 2 * p - f) // s + 1\n",
    "x = torch.zeros(1, 3, 32, 32)\n",
    "for f, p, s in ((3, 0, 1), (3, 1, 1), (5, 2, 1), (3, 1, 2), (7, 3, 2)):\n",
    "    y = nn.Conv2d(3, 8, f, stride=s, padding=p)(x)\n",
    "    print(f\"f={f} p={p} s={s}: formula {out_size(32, f, p, s)}, PyTorch {y.shape[-1]}\")\n",
    "print(\"max pool 2×2, stride 2:\", nn.MaxPool2d(2)(x).shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67cfa7ea",
   "metadata": {},
   "source": [
    "## 3. A CNN against a plain network\n",
    "Digits as 1×8×8 images. Compare accuracy and **number of parameters**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "357ac7b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.datasets import load_digits\n",
    "from sklearn.model_selection import train_test_split\n",
    "d = load_digits()\n",
    "Xa, Xb, ya, yb = train_test_split(d.images / 16.0, d.target, test_size=0.3, random_state=0)\n",
    "Xa, Xb = torch.tensor(Xa, dtype=torch.float32)[:, None], torch.tensor(Xb, dtype=torch.float32)[:, None]\n",
    "ya, yb = torch.tensor(ya), torch.tensor(yb)\n",
    "def fit(m, epochs=60, X=Xa, y=ya):\n",
    "    opt = torch.optim.AdamW(m.parameters(), lr=2e-3)\n",
    "    dl = torch.utils.data.DataLoader(torch.utils.data.TensorDataset(X, y), batch_size=64, shuffle=True)\n",
    "    for _ in range(epochs):\n",
    "        m.train()\n",
    "        for xb, yb_ in dl: opt.zero_grad(); nn.functional.cross_entropy(m(xb), yb_).backward(); opt.step()\n",
    "    m.eval()\n",
    "    with torch.no_grad(): return (m(Xb).argmax(1) == yb).float().mean().item()\n",
    "count = lambda m: sum(p.numel() for p in m.parameters())\n",
    "torch.manual_seed(0)\n",
    "plain = nn.Sequential(nn.Flatten(), nn.Linear(64, 128), nn.ReLU(), nn.Linear(128, 10))\n",
    "cnn = nn.Sequential(nn.Conv2d(1, 16, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),       # 16×4×4\n",
    "                    nn.Conv2d(16, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),      # 32×2×2\n",
    "                    nn.Flatten(), nn.Linear(128, 10))\n",
    "print(f\"plain network: {count(plain):,} parameters, test acc {fit(plain):.3f}\")\n",
    "print(f\"CNN:           {count(cnn):,} parameters, test acc {fit(cnn):.3f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cac54144",
   "metadata": {},
   "source": [
    "## 5. A residual block\n",
    "`a[l+2] = g(z[l+2] + a[l])`: the shortcut lets the block learn 'do nothing' easily, so very deep networks still train."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07e50ae2",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ResidualBlock(nn.Module):\n",
    "    def __init__(self, c):\n",
    "        super().__init__()\n",
    "        self.body = nn.Sequential(nn.Conv2d(c, c, 3, padding=1, bias=False), nn.BatchNorm2d(c), nn.ReLU(),\n",
    "                                  nn.Conv2d(c, c, 3, padding=1, bias=False), nn.BatchNorm2d(c))\n",
    "    def forward(self, x):\n",
    "        return torch.relu(self.body(x) + x)          # the skip connection\n",
    "x = torch.randn(2, 16, 8, 8)\n",
    "print(\"shape kept:\", ResidualBlock(16)(x).shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c2fa852",
   "metadata": {},
   "source": [
    "## 6. Depthwise separable convolution (MobileNet)\n",
    "The same input and output shapes, far fewer weights and multiplications."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eddcdaeb",
   "metadata": {},
   "outputs": [],
   "source": [
    "normal = nn.Conv2d(64, 128, 3, padding=1)\n",
    "separable = nn.Sequential(nn.Conv2d(64, 64, 3, padding=1, groups=64),   # depthwise: one filter per channel\n",
    "                          nn.Conv2d(64, 128, 1))                       # pointwise: 1×1 mixes channels\n",
    "x = torch.randn(1, 64, 32, 32)\n",
    "print(normal(x).shape, separable(x).shape)\n",
    "print(f\"weights: normal {count(normal):,} vs separable {count(separable):,} ({count(normal) / count(separable):.1f}× fewer)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62c03ac1",
   "metadata": {},
   "source": [
    "## 7. Data augmentation with torchvision.transforms.v2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b7ebea4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torchvision\n",
    "from torchvision.transforms import v2\n",
    "aug = v2.Compose([v2.RandomResizedCrop(8, scale=(0.7, 1.0), antialias=True), v2.RandomRotation(12),\n",
    "                  v2.ColorJitter(brightness=0.3)])\n",
    "fig, ax = plt.subplots(1, 8, figsize=(10, 1.5))\n",
    "for i, a in enumerate(ax):\n",
    "    a.imshow((Xa[0] if i == 0 else aug(Xa[0]))[0], cmap=\"gray_r\"); a.axis(\"off\"); a.set_title(\"original\" if i == 0 else \"\")\n",
    "plt.show()\n",
    "print(\"torchvision\", torchvision.__version__)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5fbd98eb",
   "metadata": {},
   "source": [
    "## 8. IoU and non-max suppression\n",
    "Our own versions, checked against `torchvision.ops`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0cce2fc5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchvision.ops import box_iou, nms\n",
    "def iou(a, b):\n",
    "    x1, y1 = max(a[0], b[0]), max(a[1], b[1]); x2, y2 = min(a[2], b[2]), min(a[3], b[3])\n",
    "    inter = max(0, x2 - x1) * max(0, y2 - y1)\n",
    "    return inter / ((a[2] - a[0]) * (a[3] - a[1]) + (b[2] - b[0]) * (b[3] - b[1]) - inter)\n",
    "def my_nms(boxes, scores, thr=0.5):\n",
    "    order, keep = list(np.argsort(-scores)), []\n",
    "    while order:\n",
    "        i = order.pop(0); keep.append(int(i))\n",
    "        order = [j for j in order if iou(boxes[i], boxes[j]) <= thr]\n",
    "    return keep\n",
    "boxes = np.array([[30, 30, 130, 120], [40, 38, 140, 128], [35, 25, 125, 125], [200, 60, 290, 160], [210, 70, 300, 165]], float)\n",
    "scores = np.array([0.9, 0.75, 0.6, 0.85, 0.7])\n",
    "print(\"IoU:\", round(iou(boxes[0], boxes[1]), 3), \"torchvision:\", round(box_iou(torch.tensor(boxes[:1]), torch.tensor(boxes[1:2])).item(), 3))\n",
    "print(\"NMS keeps:\", my_nms(boxes, scores), \"torchvision:\", nms(torch.tensor(boxes), torch.tensor(scores), 0.5).tolist())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4144f5a7",
   "metadata": {},
   "source": [
    "## 9. Transfer learning with pre-trained weights\n",
    "The modern torchvision API: `weights=...DEFAULT` and the matching `transforms()`. Freeze the backbone, replace the last layer. (Downloads about 45 MB, so it needs internet: it runs in Colab.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e159d781",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchvision.models import resnet18, ResNet18_Weights\n",
    "try:\n",
    "    weights = ResNet18_Weights.DEFAULT\n",
    "    net = resnet18(weights=weights)\n",
    "    for p in net.parameters(): p.requires_grad = False\n",
    "    net.fc = nn.Linear(net.fc.in_features, 3)            # your 3 classes; only this layer trains\n",
    "    print(\"trainable parameters:\", sum(p.numel() for p in net.parameters() if p.requires_grad))\n",
    "    print(\"preprocessing:\", weights.transforms())\n",
    "except Exception as e:\n",
    "    print(\"Needs internet to download the weights (run this in Colab):\", type(e).__name__)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db6bc723",
   "metadata": {},
   "source": [
    "## 10. Triplet loss and the style matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99344354",
   "metadata": {},
   "outputs": [],
   "source": [
    "def triplet_loss(a, p, n, margin=0.2):\n",
    "    return torch.relu((a - p).pow(2).sum(1) - (a - n).pow(2).sum(1) + margin).mean()\n",
    "a = nn.functional.normalize(torch.randn(4, 128), dim=1)          # embeddings on the unit sphere, like FaceNet\n",
    "pos = nn.functional.normalize(a + 0.3 * torch.randn_like(a), dim=1)  # same person: close\n",
    "neg = nn.functional.normalize(torch.randn(4, 128), dim=1)            # someone else: far\n",
    "hard = nn.functional.normalize(a + 0.35 * torch.randn_like(a), dim=1)  # a \"hard\" negative: almost as close\n",
    "print(\"easy negatives:\", round(triplet_loss(a, pos, neg).item(), 3), \"| hard negatives:\", round(triplet_loss(a, pos, hard).item(), 3))\n",
    "print(\"PyTorch's built-in (plain distances):\", round(nn.TripletMarginLoss(margin=0.2)(a, pos, hard).item(), 3))\n",
    "def gram(features):                     # features: (channels, height, width)\n",
    "    f = features.flatten(1)\n",
    "    return f @ f.T / f.shape[1]          # which channels fire together = the \"style\"\n",
    "print(\"style matrix shape for 64 channels:\", gram(torch.randn(64, 14, 14)).shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "910688d1",
   "metadata": {},
   "source": [
    "## Try this\n",
    "1. Remove the pooling layers in section 3 (and fix the `Linear` size). More parameters, better accuracy?\n",
    "2. Train the CNN on 20% of the training data, with and without the augmentation from section 7.\n",
    "3. In section 8, change the NMS threshold to 0.3 and 0.7. Which boxes survive?\n",
    "4. In Colab, fine-tune section 9's ResNet on a small image folder of your own (`torchvision.datasets.ImageFolder`)."
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}