{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lesson 36: Transfer Learning\n",
    "\n",
    "Training a good CNN from scratch (Lessons 33-34) took hundreds of labeled examples even for a toy task. Real target tasks are often data-starved: a handful of labeled medical scans, a new product category with 20 photos. **Transfer learning** sidesteps this by reusing a network already trained on a *different*, data-rich task, on the theory that early-layer features (edges, blobs, simple textures — Lesson 33's Sobel-like first-layer filters) are useful for almost any visual task, not just the one they were originally trained on."
   ]
  },
  {
   "cell_type": "code",
   "id": "f2c4e96e",
   "source": "import copy\nimport numpy as np\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport matplotlib.pyplot as plt",
   "metadata": {},
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "2112ff86",
   "source": "## Source task and target task\n\nSet up two related but distinct tasks. The **source task** has plenty of data: distinguishing plusses from circles, 300 training images. The **target task** is the one we actually care about, and it's deliberately starved: distinguishing squares from circles (a new class the source task never saw), with only **12 noisy training images**.",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "id": "2c0bfbc4",
   "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    elif shape_type == 'circle':\n        yy, xx = np.mgrid[0:size, 0:size]\n        img[((xx-cx)**2 + (yy-cy)**2) <= 9] = 1.0\n    elif shape_type == 'square':\n        img[cy-3:cy+4, cx-3:cx+4] = 1.0\n    return img\n\ndef make_dataset(rng_local, n, shapes, position_range=(5, 11), noise=0.15):\n    imgs, labels = [], []\n    for _ in range(n):\n        shape_type = rng_local.choice(shapes)\n        cx, cy = rng_local.integers(*position_range), rng_local.integers(*position_range)\n        img = make_image(shape_type, cx, cy)\n        img = np.clip(img + rng_local.normal(0, noise, img.shape), 0, 1).astype(np.float32)\n        imgs.append(img)\n        labels.append(shapes.index(shape_type))\n    return np.array(imgs, dtype=np.float32), np.array(labels, dtype=np.int64)\n\ndata_rng = np.random.default_rng(3)\nX_src, y_src = make_dataset(data_rng, 300, ['plus', 'circle'], noise=0.15)\nX_tgt_train, y_tgt_train = make_dataset(data_rng, 12, ['square', 'circle'], noise=0.4)\nX_tgt_test, y_tgt_test = make_dataset(data_rng, 150, ['square', 'circle'], noise=0.4)\n\nfig, axes = plt.subplots(2, 6, figsize=(11, 4))\nfor ax, im in zip(axes[0], X_src[:6]):\n    ax.imshow(im, cmap='gray'); ax.axis('off')\naxes[0, 0].set_ylabel('source', rotation=0, labelpad=25)\nfor ax, im in zip(axes[1], X_tgt_train[:6]):\n    ax.imshow(im, cmap='gray'); ax.axis('off')\naxes[1, 0].set_ylabel('target', rotation=0, labelpad=25)\nfig.suptitle('Source task (plus vs. circle, top) vs. target task (square vs. circle, bottom)', y=1.02)\nplt.show()",
   "metadata": {},
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "b9c8ff23",
   "source": "## Three strategies\n\n1. **From scratch** — train a fresh CNN on only the 12 target images. This is the baseline: no transfer at all.\n2. **Frozen backbone (feature extraction)** — pretrain a CNN backbone on the source task, then freeze its weights entirely and train only a new linear classifier on top of the features it produces for target images.\n3. **Fine-tuning** — start from the same pretrained backbone, but keep updating it on the target data too, using a *much smaller* learning rate for the backbone than for the new classifier head (the backbone already encodes useful structure; large updates from just 12 examples would wreck it).",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "id": "a183e046",
   "source": "class Backbone(nn.Module):\n    def __init__(self):\n        super().__init__()\n        self.conv = nn.Sequential(\n            nn.Conv2d(1, 16, 5, padding=2), nn.ReLU(),\n            nn.MaxPool2d(2),\n            nn.Conv2d(16, 32, 5, padding=2), nn.ReLU(),\n            nn.AdaptiveMaxPool2d(1),\n        )\n\n    def forward(self, x):\n        return self.conv(x).flatten(1)\n\nclass Classifier(nn.Module):\n    def __init__(self, backbone, freeze_backbone):\n        super().__init__()\n        self.backbone = backbone\n        self.freeze_backbone = freeze_backbone\n        self.fc = nn.Linear(32, 2)\n\n    def forward(self, x):\n        feat = self.backbone(x)\n        if self.freeze_backbone:\n            feat = feat.detach()  # no gradient flows into a frozen backbone\n        return self.fc(feat)\n\ndef train_source(seed, epochs=300, lr=0.01):\n    torch.manual_seed(seed)  # seed before constructing the model (Lesson 33/34)\n    backbone = Backbone()\n    fc = nn.Linear(32, 2)\n    opt = torch.optim.Adam(list(backbone.parameters()) + list(fc.parameters()), lr=lr)\n    Xt = torch.tensor(X_src).unsqueeze(1); yt = torch.tensor(y_src)\n    for _ in range(epochs):\n        opt.zero_grad()\n        loss = F.cross_entropy(fc(backbone(Xt)), yt)\n        loss.backward()\n        opt.step()\n    return backbone\n\ndef train_target(backbone, freeze_backbone, seed, epochs=200, lr=0.01, backbone_lr=None):\n    torch.manual_seed(seed)\n    model = Classifier(backbone, freeze_backbone)\n    if freeze_backbone:\n        opt = torch.optim.Adam(model.fc.parameters(), lr=lr)\n    elif backbone_lr is not None:\n        opt = torch.optim.Adam([\n            {'params': model.backbone.parameters(), 'lr': backbone_lr},\n            {'params': model.fc.parameters(), 'lr': lr},\n        ])\n    else:\n        opt = torch.optim.Adam(model.parameters(), lr=lr)\n    Xtr = torch.tensor(X_tgt_train).unsqueeze(1); ytr = torch.tensor(y_tgt_train)\n    Xte = torch.tensor(X_tgt_test).unsqueeze(1); yte = torch.tensor(y_tgt_test)\n    for _ in range(epochs):\n        opt.zero_grad()\n        loss = F.cross_entropy(model(Xtr), ytr)\n        loss.backward()\n        opt.step()\n    with torch.no_grad():\n        return (model(Xte).argmax(1) == yte).float().mean().item()",
   "metadata": {},
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "id": "23e5da56",
   "source": "scratch_accs, frozen_accs, finetune_accs = [], [], []\nfor seed in range(8):\n    scratch_acc = train_target(Backbone(), freeze_backbone=False, seed=seed)\n    pretrained = train_source(seed=seed)\n    frozen_acc = train_target(copy.deepcopy(pretrained), freeze_backbone=True, seed=seed)\n    finetune_acc = train_target(copy.deepcopy(pretrained), freeze_backbone=False, seed=seed, backbone_lr=0.0005)\n    scratch_accs.append(scratch_acc); frozen_accs.append(frozen_acc); finetune_accs.append(finetune_acc)\n\nprint(f'{\"strategy\":>20} {\"mean test acc\":>16} {\"std\":>8}')\nprint(f'{\"from scratch\":>20} {np.mean(scratch_accs):>15.1%} {np.std(scratch_accs):>8.1%}')\nprint(f'{\"frozen backbone\":>20} {np.mean(frozen_accs):>15.1%} {np.std(frozen_accs):>8.1%}')\nprint(f'{\"fine-tuned\":>20} {np.mean(finetune_accs):>15.1%} {np.std(finetune_accs):>8.1%}')\nprint(f'\\n(averaged over {len(scratch_accs)} random seeds, each reusing the same 12 target training images)')",
   "metadata": {},
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "78b1a292",
   "source": "Both transfer strategies beat training from scratch on average, and both are also noticeably more *consistent* (lower standard deviation) — with only 12 training images, a from-scratch network's success or failure depends heavily on which 12 images it happened to get, while a pretrained backbone starts from a much better place regardless. The source task never saw a square, yet the low-level features it learned (edges, curvature, blob-like regions) transferred anyway, because those features are generic to shape recognition, not specific to \"plus vs. circle.\"\n\nNote what fine-tuning needed to work at all: a backbone learning rate roughly 20x smaller than the classifier head's. With only 12 examples, an unrestrained backbone update would simply overfit those 12 images from scratch, discarding everything useful it learned from the 300-image source task — the same catastrophic-forgetting failure mode as Lesson 34's overfitting curve, just applied to a network that started out already knowing something.\n\n### Exercise\n\n1. Try `backbone_lr=0.01` (i.e. no learning-rate difference between backbone and head) in the fine-tuning call. Does fine-tuned accuracy get better or worse, and does that match the catastrophic-forgetting explanation above?\n2. Try freezing *most* of the backbone but fine-tuning only its last conv layer (hint: set `requires_grad = False` on the first `Conv2d`'s parameters only). Where does that land relative to fully-frozen and fully-fine-tuned?\n3. The source and target tasks here share a class (circle). Design a source task that shares *no* classes with the target task at all, and predict whether transfer would still help. Test your prediction.",
   "metadata": {}
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.10.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}