{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3ed9cb4a",
   "metadata": {},
   "source": "# Lesson 50: Self-Supervised and Contrastive Learning\n\nLesson 38'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 {{lesson:geometric_transformations}}/{{lesson:convolutional_neural_networks}}'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": "444a3027",
   "source": "## A small Vision Transformer as the encoder\n\nThe encoder that turns each augmented view into an embedding is a ViT (Lesson {{lesson:vision_transformers}}'s exact pattern: patchify, linearly embed, a `[CLS]` token gathers information via self-attention, project its final state to the embedding dimension) rather than a CNN. Real SimCLR's own paper (<a href=\"../references.html#chen-2020-simclr\">Chen et al., 2020</a>) used a CNN (ResNet) backbone instead — also a perfectly legitimate, arguably more sample-efficient choice at this scale — but this part of the course is specifically about Transformer-based architectures, so a ViT is used here instead.",
   "metadata": {}
  },
  {
   "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": "def patchify(img, patch_size):\n    B, C, H, W = img.shape\n    P = patch_size\n    patches = img.unfold(2, P, P).unfold(3, P, P)  # (B, C, H/P, W/P, P, P)\n    patches = patches.contiguous().view(B, C, -1, P, P).permute(0, 2, 1, 3, 4)\n    return patches.reshape(B, -1, C * P * P)  # (B, num_patches, C*P*P)\n\ndef positional_encoding(T, D):\n    pos = torch.arange(T).unsqueeze(1).float()\n    i = torch.arange(D).unsqueeze(0).float()\n    angle_rates = 1.0 / (10000 ** (2 * (i // 2) / D))\n    angles = pos * angle_rates\n    pe = torch.zeros(T, D)\n    pe[:, 0::2] = torch.sin(angles[:, 0::2])\n    pe[:, 1::2] = torch.cos(angles[:, 1::2])\n    return pe\n\nclass Encoder(nn.Module):\n    def __init__(self, out_dim=16, img_size=16, patch_size=4, embed_dim=32, n_heads=4, n_layers=2):\n        super().__init__()\n        self.patch_size = patch_size\n        n_patches = (img_size // patch_size) ** 2\n        patch_dim = patch_size * patch_size\n        self.embed = nn.Linear(patch_dim, embed_dim)\n        self.cls_token = nn.Parameter(torch.randn(1, 1, embed_dim) * 0.02)\n        self.register_buffer('pos_enc', positional_encoding(n_patches + 1, embed_dim))\n        layer = nn.TransformerEncoderLayer(embed_dim, n_heads, dim_feedforward=embed_dim * 2,\n                                            batch_first=True, dropout=0.0)\n        self.encoder = nn.TransformerEncoder(layer, num_layers=n_layers)\n        self.proj = nn.Linear(embed_dim, out_dim)\n\n    def forward(self, x):\n        p = patchify(x, self.patch_size)\n        tok = self.embed(p)\n        cls = self.cls_token.expand(x.shape[0], -1, -1)\n        tok = torch.cat([cls, tok], dim=1) + self.pos_enc\n        out = self.encoder(tok)\n        return self.proj(out[:, 0])  # project the [CLS] token's final state to out_dim\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 {{lesson:transfer_learning}}'s frozen-backbone pattern) and train only a small linear classifier on top of its features — a technique called **linear probing** — 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 shows a modest but consistently positive improvement over random features across repeated runs — not the dramatic gap a CNN encoder shows on this same setup, but a real one. That's an honest consequence of the architecture choice, not a bug: Lesson {{lesson:vision_transformers}} already showed that a ViT has no built-in translation invariance and needs far more data than a CNN to learn useful structure from scratch, and the same limitation shows up here — the contrastive objective still pushes same-shape images together and different-shape images apart, but a few hundred unlabeled toy images and a couple hundred training steps is a lot less signal for a Transformer to work with than it would be for a CNN's stronger inductive bias. (This is also why SimCLR's own original paper used a CNN/ResNet backbone, not a ViT — a legitimate, arguably more sample-efficient choice at this scale; the ViT is used here instead because this part of the course is specifically about Transformer-based self-supervised pretraining.)\n\nThis is the mechanism behind why self-supervised pretraining matters at scale: models like DINOv2 (Lesson {{lesson:dino_self_distillation}}) 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. At real scale, the data-hunger gap this lesson's toy experiment exposes stops being a limitation at all.\n\n### Exercises\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
}