{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "# Lesson 47: Self-Supervised and Contrastive Learning\n\nLesson 36's transfer learning still needed a labeled source task to pretrain a useful backbone. **Self-supervised learning** removes even that requirement: pretrain on *unlabeled* images by inventing a task from the data itself, using no human annotations at all. This lesson builds the dominant recipe for that, **contrastive learning** (SimCLR-style, <a href=\"../references.html#chen-2020-simclr\">Chen et al., 2020</a>): two randomly augmented views of the same image should produce similar embeddings; views of different images should not."
  },
  {
   "cell_type": "code",
   "id": "2e3a8417",
   "source": "import numpy as np\nimport cv2\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": "7217b778",
   "source": "## Two views, one identity\n\nTake 500 unlabeled shape images (the labels exist in this synthetic dataset only so accuracy can be measured later — the pretraining step below never looks at them). For each image sampled during pretraining, create two independently augmented views (Lesson 8/34's flips and rotations, plus pixel noise).",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "id": "a6b9347d",
   "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\ndef make_dataset(rng_local, n, position_range=(4, 12)):\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 if shape_type == 'plus' else 1)\n    return np.array(imgs, dtype=np.float32), np.array(labels, dtype=np.int64)\n\ndef augment(img, rng_local):\n    if rng_local.random() < 0.5:\n        img = np.fliplr(img).copy()\n    angle = rng_local.uniform(-20, 20)\n    M = cv2.getRotationMatrix2D((8, 8), angle, 1.0)\n    img = cv2.warpAffine(img, M, (16, 16))\n    return np.clip(img + rng_local.normal(0, 0.1, img.shape), 0, 1).astype(np.float32)\n\nrng = np.random.default_rng(5)\nX_unlabeled, y_unlabeled = make_dataset(rng, 500)\n\naug_rng = np.random.default_rng(1)\nfig, axes = plt.subplots(2, 4, figsize=(9, 4.5))\nfor i in range(4):\n    axes[0, i].imshow(augment(X_unlabeled[i], aug_rng), cmap='gray'); axes[0, i].axis('off')\n    axes[1, i].imshow(augment(X_unlabeled[i], aug_rng), cmap='gray'); axes[1, i].axis('off')\naxes[0, 0].set_title('view 1', fontsize=9, loc='left')\naxes[1, 0].set_title('view 2', fontsize=9, loc='left')\nplt.suptitle('Two augmented views of the same 4 images')\nplt.show()",
   "metadata": {},
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "8e40ab55",
   "source": "## The contrastive loss (NT-XEnt)\n\nFor a batch of `N` images, produce `2N` embeddings (two views each). For each embedding, its one true positive is the *other* view of the same image; every other one of the `2N-2` embeddings is a negative. Treat this as a classification problem — \"which of the other 2N-1 embeddings is my positive pair?\" — and minimize cross-entropy over cosine similarities. This is the **normalized temperature-scaled cross-entropy (NT-Xent)** loss.",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "id": "d609aa38",
   "source": "class Encoder(nn.Module):\n    def __init__(self, out_dim=16):\n        super().__init__()\n        self.conv = nn.Sequential(\n            nn.Conv2d(1, 16, 5, padding=2), nn.ReLU(), nn.MaxPool2d(2),\n            nn.Conv2d(16, 32, 5, padding=2), nn.ReLU(), nn.AdaptiveMaxPool2d(1),\n        )\n        self.proj = nn.Linear(32, out_dim)\n\n    def forward(self, x):\n        return self.proj(self.conv(x).flatten(1))\n\ndef nt_xent_loss(z1, z2, temperature=0.5):\n    z1 = F.normalize(z1, dim=1)\n    z2 = F.normalize(z2, dim=1)\n    z = torch.cat([z1, z2], dim=0)          # (2N, D)\n    sim = z @ z.T / temperature              # (2N, 2N) cosine similarities\n    N = z1.shape[0]\n    mask = torch.eye(2 * N, dtype=torch.bool)\n    sim.masked_fill_(mask, float('-inf'))    # exclude comparing an embedding to itself\n    targets = torch.cat([torch.arange(N, 2 * N), torch.arange(0, N)])  # each row's true positive index\n    return F.cross_entropy(sim, targets)\n\ndef train_contrastive(seed, epochs=200, lr=0.01, batch_size=64):\n    torch.manual_seed(seed)\n    encoder = Encoder()\n    opt = torch.optim.Adam(encoder.parameters(), lr=lr)\n    local_aug_rng = np.random.default_rng(seed + 100)\n    n = len(X_unlabeled)\n    for epoch in range(epochs):\n        idx = np.random.default_rng(epoch).permutation(n)[:batch_size]\n        batch = X_unlabeled[idx]\n        view1 = np.stack([augment(im, local_aug_rng) for im in batch])\n        view2 = np.stack([augment(im, local_aug_rng) for im in batch])\n        z1 = encoder(torch.tensor(view1).unsqueeze(1))\n        z2 = encoder(torch.tensor(view2).unsqueeze(1))\n        loss = nt_xent_loss(z1, z2)\n        opt.zero_grad()\n        loss.backward()\n        opt.step()\n    return encoder, loss.item()\n\nencoder, final_loss = train_contrastive(seed=0)\nprint(f'final NT-Xent loss after 200 steps: {final_loss:.3f}  (random-chance loss would be ln(2*64-1) = {np.log(2*64-1):.3f})')",
   "metadata": {},
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "e9142da6",
   "source": "## Evaluating with a linear probe\n\nThe standard way to check whether self-supervised pretraining actually learned something useful: freeze the pretrained encoder (Lesson 36's frozen-backbone pattern) and train only a small linear classifier on top of its features, using a *tiny* number of labeled examples — exactly the low-label regime self-supervised pretraining is meant to help with. Compare against the same linear probe on a randomly initialized (never trained) encoder.",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "id": "421b2a1b",
   "source": "def make_noisy_dataset(rng_local, n, noise=0.35, position_range=(4, 12)):\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        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(0 if shape_type == 'plus' else 1)\n    return np.array(imgs, dtype=np.float32), np.array(labels, dtype=np.int64)\n\nX_probe_train, y_probe_train = make_noisy_dataset(rng, 8)    # only 8 labeled examples\nX_probe_test, y_probe_test = make_noisy_dataset(rng, 150)\n\ndef linear_probe_acc(enc, seed):\n    torch.manual_seed(seed)\n    with torch.no_grad():\n        feat_train = enc(torch.tensor(X_probe_train).unsqueeze(1))\n        feat_test = enc(torch.tensor(X_probe_test).unsqueeze(1))\n    probe = nn.Linear(feat_train.shape[1], 2)\n    opt = torch.optim.Adam(probe.parameters(), lr=0.05)\n    ytr = torch.tensor(y_probe_train)\n    for _ in range(300):\n        opt.zero_grad()\n        loss = F.cross_entropy(probe(feat_train), ytr)\n        loss.backward()\n        opt.step()\n    with torch.no_grad():\n        preds = probe(feat_test).argmax(1).numpy()\n    return (preds == y_probe_test).mean()\n\nrandom_accs, contrastive_accs = [], []\nfor seed in range(5):\n    enc_contrastive, _ = train_contrastive(seed=seed)\n    contrastive_accs.append(linear_probe_acc(enc_contrastive, seed=seed + 50))\n\n    torch.manual_seed(seed)\n    enc_random = Encoder()\n    random_accs.append(linear_probe_acc(enc_random, seed=seed + 50))\n\nprint(f'linear probe on RANDOM (untrained) features:     {np.mean(random_accs):.1%}  (+/- {np.std(random_accs):.1%})')\nprint(f'linear probe on CONTRASTIVE-pretrained features: {np.mean(contrastive_accs):.1%}  (+/- {np.std(contrastive_accs):.1%})')",
   "metadata": {},
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "7be0ec61",
   "source": "With only 8 labeled examples to train the probe, the contrastively-pretrained encoder gives a real (if noisy, given how few labels are involved) improvement over random features — the contrastive objective, despite never seeing a label, has pushed same-shape images toward nearby points in embedding space and different-shape images apart, purely by requiring that augmented views of the same image agree.\n\nThis is the mechanism behind why self-supervised pretraining matters at scale: models like DINOv2 (Lesson 48) are pretrained this way on hundreds of millions of *unlabeled* images — impossible to hand-label at that scale — and the resulting features transfer to downstream tasks with only a handful of labeled examples, exactly as demonstrated here in miniature.\n\n### Exercise\n\n1. Increase `out_dim` in `Encoder` from 16 to 64. Does a higher-dimensional embedding space improve the linear probe's accuracy, hurt it, or make little difference at this data scale?\n2. Try `temperature=0.1` and `temperature=2.0` in `nt_xent_loss` instead of `0.5`. Temperature controls how sharply the loss penalizes near-miss negatives — does either extreme change the final probe accuracy noticeably?\n3. Weaken the augmentations in `augment` (e.g. rotation range `(-5, 5)` instead of `(-20, 20)`, no noise). Does contrastive pretraining still beat the random baseline — and what does the answer suggest about *why* augmentation strength is one of the most-tuned hyperparameters in real contrastive learning pipelines?",
   "metadata": {}
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.10.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}