Lesson 112: Instance Segmentation & Contrastive Learning

1. Instance Segmentation & Contrastive Learning

Title

USAAIO · Lesson 112 · Phase 3

Mask R-CNN segments every object instance pixel-by-pixel. SimCLR learns visual representations without a single label by pulling augmented views together and pushing different images apart. Both topics verified in PyTorch with real computed shapes and losses.

2. By the end of this lesson you can

Objectives

1. Explain how Mask R-CNN extends Faster R-CNN with a mask head and trace RoI feature shapes through it
2. Distinguish instance segmentation from semantic segmentation and from panoptic segmentation
3. Write the InfoNCE / NT-Xent loss from scratch and explain the role of temperature tau
4. Describe the SimCLR pipeline: augmentation → encoder → projection head → contrastive loss
5. Evaluate self-supervised representations via a linear probe and compare to supervised baselines

3. Mask R-CNN — instance segmentation

Section

Part 1 of 3

4. Instance vs semantic vs panoptic segmentation

Concept

Three pixel-level vision tasks are often confused. The key axis: does the model count and separate individual objects, or assign each pixel a category label?

Things (countable objects: cars, people) get per-instance masks. Stuff (amorphous regions: sky, road) only gets a semantic label. Panoptic = things + stuff.

5. Fill in: assigns label per pixel for Instance vs semantic vs panoptic segmentation

Comparison

6. Mask R-CNN: Faster R-CNN + mask head

Concept

Mask R-CNN (He 2017) takes Faster R-CNN (Lesson 56) and adds a third head in parallel with the box and class heads. The mask head operates on RoI-aligned features — not RoI-pooled — which matters for pixel-level accuracy.

• Backbone + FPN: ResNet-50 or 101 + feature pyramid (same as Faster R-CNN)
• RPN: proposes regions of interest (same as Faster R-CNN)
• RoIAlign: bilinear interpolation — avoids the quantization error of RoIPool; output (N, C, 7, 7) for box/class head, (N, C, 14, 14) for mask head
• Mask head: small conv stack + ConvTranspose2d → (N, K, 28, 28) binary mask per class K

7. Guess the shape of the answer: Mask head — shape trace in PyTorch

Estimation

8. Mask head — shape trace in PyTorch

Worked example

Trace a minimal mask head: 4 proposals, 256-channel RoI features at 14×14, 5 classes. Follow the shapes through each layer before reading on.

import torch, torch.nn as nn

class TinyMaskHead(nn.Module):
    def __init__(self, in_ch=256, n_cls=5):
        super().__init__()
        self.convs = nn.Sequential(
            nn.Conv2d(in_ch, 64, 3, padding=1), nn.ReLU(),
            nn.Conv2d(64,  64, 3, padding=1), nn.ReLU(),
        )
        self.deconv = nn.ConvTranspose2d(64, 64, 2, stride=2)  # 14->28
        self.pred   = nn.Conv2d(64, n_cls, 1)                   # 1x1 per-class
    def forward(self, x):
        x = self.convs(x)
        x = self.deconv(x)
        return self.pred(x)

torch.manual_seed(42)
head = TinyMaskHead(in_ch=256, n_cls=5)
roi  = torch.randn(4, 256, 14, 14)   # 4 proposals
out  = head(roi)
print(roi.shape)  # torch.Size([4, 256, 14, 14])
print(out.shape)  # torch.Size([4, 5, 28, 28])
print(sum(p.numel() for p in head.parameters()))  # 201221

9. What each one costs: Mask head — shape trace in PyTorch

Trade off

10. Panoptic segmentation

Concept

Panoptic segmentation (Kirillov 2019) assigns every pixel either a thing label with an instance ID, or a stuff label with no instance ID. The Panoptic Quality (PQ) metric unifies detection and segmentation quality:

\[ \text{PQ} = \underbrace{\frac{\sum_{(p,g)\in TP} \text{IoU}(p,g)}{|TP|}}_{{\text{Segmentation Quality (SQ)}}} \times \underbrace{\frac{|TP|}{|TP| + \tfrac{1}{2}|FP| + \tfrac{1}{2}|FN|}}_{{\text{Recognition Quality (RQ)}}} \]

A modern approach is a single unified model (e.g. Panoptic FPN, Mask2Former) that shares a backbone and FPN between the semantic and instance branches, merging outputs in post-processing.

11. Break it if you can: Panoptic segmentation

Counterexample

12. Something is wrong here: RoIPool vs RoIAlign in Mask R-CNN

Anomaly

13. Trap: RoIPool vs RoIAlign in Mask R-CNN

Trap

14. Break it on purpose: RoIPool vs RoIAlign in Mask R-CNN

Break the constraint

15. Contrastive learning — SimCLR & InfoNCE

Section

Part 2 of 3

16. Self-supervised: learning without labels

Concept

Supervised pretraining needs millions of labeled images. Self-supervised pretraining creates a proxy task from the data itself — no labels required. The learned encoder can then be fine-tuned on a small labeled set.

Contrastive learning formulates the proxy task as: given two augmented views of the same image, their embeddings should be similar; embeddings of different images should be dissimilar. SimCLR (Chen 2020) is the canonical instantiation.

• Positive pair: two random augmented crops of the same image
• Negatives: augmented views of all other images in the batch
• Signal: no labels needed — the data augmentation defines what 'same' means

17. InfoNCE loss — the math

Concept

The InfoNCE loss (Oord 2018) treats contrastive learning as a 2N-1-way classification problem: given an anchor z_i, identify its positive z_j among all 2(N-1) negatives in the batch.

\[ \mathcal{L}_{i,j} = -\log \frac{\exp(\text{sim}(z_i, z_j) / \tau)}{\sum_{k=1, k \neq i}^{2N} \exp(\text{sim}(z_i, z_k) / \tau)} \]

• sim(u, v) = uᵀv / (‖u‖ · ‖v‖) — cosine similarity after L2-normalizing both
• tau (temperature): small tau sharpens the distribution, making the loss harder; large tau is softer
• NT-Xent (normalized temperature cross-entropy) = SimCLR's symmetric InfoNCE: loss is averaged over both (i,j) and (j,i) directions

18. By analogy: InfoNCE loss — the math

Analogy

19. What has to be given first: InfoNCE from scratch — code + trace

Missing information

20. InfoNCE from scratch — code + trace

Worked example

Implement NT-Xent (symmetric InfoNCE) for a batch of B=6 pairs. Both views are stacked into a 2B×D matrix; the similarity matrix is (2B, 2B). Verify the loss numerically.

import torch, torch.nn as nn

def nt_xent_loss(z1, z2, tau=0.5):
    """z1, z2: (B, D) embeddings from two augmented views."""
    B = z1.shape[0]
    z1 = nn.functional.normalize(z1, dim=-1)
    z2 = nn.functional.normalize(z2, dim=-1)
    z  = torch.cat([z1, z2], dim=0)      # (2B, D)
    sim = z @ z.T / tau                   # (2B, 2B) cosine / tau
    # Mask self-similarities (diagonal = same vector)
    mask = torch.eye(2*B, dtype=torch.bool)
    sim.masked_fill_(mask, -1e9)
    # For index i in [0..B-1]: positive is i+B; for [B..2B-1]: positive is i-B
    labels = torch.cat([torch.arange(B, 2*B), torch.arange(B)])
    return nn.CrossEntropyLoss()(sim, labels)

torch.manual_seed(42)
B, D = 6, 8
z1 = torch.randn(B, D)
z2 = z1 + 0.05 * torch.randn_like(z1)  # tight positives
print(f'tau=0.07: {nt_xent_loss(z1, z2, 0.07):.4f}')   # 0.0009
print(f'tau=0.50: {nt_xent_loss(z1, z2, 0.50):.4f}')   # 0.8428
print(f'tau=1.00: {nt_xent_loss(z1, z2, 1.00):.4f}')   # 1.4959

21. Work backwards from the answer: InfoNCE from scratch — code + trace

Reverse engineer

22. SimCLR pipeline: augment → encode → project → loss

Concept

SimCLR has four components. Critically, the projection head g(·) is used only during pretraining — the linear probe and fine-tuning attach to the encoder output h, not z.

Why discard g? Chen 2020 Appendix B shows that representations in h transfer better than in z — the projection head learns to discard information (e.g. color, orientation) that the augmentations remove but that is useful for downstream tasks.

23. By analogy: SimCLR pipeline: augment → encode → project → loss

Analogy

24. Something is wrong here: fine-tuning from the projection head output z

Anomaly

25. Trap: fine-tuning from the projection head output z

Trap

26. Which of these survive contact with Lesson 112: Instance Segmentation &…?

Two truths and a lie

27. Guess the shape of the answer: SimCLR architecture — encoder + projection…

Estimation

28. SimCLR architecture — encoder + projection head

Worked example

Build a minimal SimCLR (encoder + projection head) on synthetic 64-dim inputs. Verify parameter counts and embedding shapes before running.

import torch, torch.nn as nn

class TinyEncoder(nn.Module):
    def __init__(self, in_dim=64, hidden=32, out_dim=16):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(in_dim, hidden), nn.ReLU(),
            nn.Linear(hidden, out_dim)
        )
    def forward(self, x): return self.net(x)

class ProjectionHead(nn.Module):
    def __init__(self, in_dim=16, hidden=16, out_dim=8):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(in_dim, hidden), nn.ReLU(),
            nn.Linear(hidden, out_dim)
        )
    def forward(self, x): return self.net(x)

torch.manual_seed(42)
enc  = TinyEncoder(64, 32, 16)
proj = ProjectionHead(16, 16, 8)
print(sum(p.numel() for p in enc.parameters()),   # 2608
      sum(p.numel() for p in proj.parameters()))  # 408

x   = torch.randn(8, 64)
h   = enc(x)              # (8, 16) - representation
z   = proj(h)             # (8, 8)  - projection
print(h.shape, z.shape)   # torch.Size([8, 16]) torch.Size([8, 8])

29. Work backwards from the answer: SimCLR architecture — encoder + projection…

Reverse engineer

30. Linear probe & self-supervised evaluation

Section

Part 3 of 3

31. Linear probe: measuring representation quality

Concept

A linear probe measures self-supervised representation quality: freeze the pretrained encoder, train only a linear classifier on top of h. If a single linear layer achieves high accuracy, the encoder has learned linearly separable representations.

• Why linear? A linear probe isolates the encoder — it cannot compensate for bad representations with extra capacity
• Evaluation protocol: train the linear head on the labeled set; measure top-1/top-5 accuracy on a held-out test set
• SimCLR result (ImageNet): 76.5% top-1 with a linear probe vs 76.3% for full supervised ResNet-50 — self-supervised nearly matches supervised at the same model size

32. Guess the shape of the answer: SimCLR pretraining + linear probe on…

Estimation

33. SimCLR pretraining + linear probe on synthetic data

Worked example

Run SimCLR on 160 unlabeled synthetic 64-dim samples (two classes). Then attach a frozen-encoder linear probe and compare to a supervised baseline. Predict whether SSL or supervised wins.

import torch, torch.nn as nn

# --- Data: two clusters, no labels during pretraining ---
torch.manual_seed(42)
vec0 = torch.zeros(64); vec0[0] = 2.0
vec1 = torch.zeros(64); vec1[0] = -2.0
X_all = torch.cat([torch.randn(100,64)+vec0, torch.randn(100,64)+vec1])
y_all = torch.cat([torch.zeros(100,dtype=torch.long), torch.ones(100,dtype=torch.long)])
perm  = torch.randperm(200, generator=torch.Generator().manual_seed(42))
X_tr, X_te = X_all[perm[:160]], X_all[perm[160:]]
y_tr, y_te = y_all[perm[:160]], y_all[perm[160:]]

# --- SSL pretraining (no y_tr used) ---
torch.manual_seed(42)
encoder  = TinyEncoder(64, 32, 16)   # defined earlier
projector= ProjectionHead(16, 16, 8)
opt_ssl  = torch.optim.Adam(list(encoder.parameters())+list(projector.parameters()), lr=1e-3)
for epoch in range(60):
    g1=torch.Generator().manual_seed(epoch)
    g2=torch.Generator().manual_seed(epoch+100)
    v1 = encoder(X_tr + 0.1*torch.randn(X_tr.shape, generator=g1))
    v2 = encoder(X_tr + 0.1*torch.randn(X_tr.shape, generator=g2))
    loss = nt_xent_loss(projector(v1), projector(v2), tau=0.5)  # defined earlier
    opt_ssl.zero_grad(); loss.backward(); opt_ssl.step()

# --- Linear probe ---
with torch.no_grad():
    h_tr = encoder(X_tr); h_te = encoder(X_te)
lin = nn.Linear(16, 2)
opt_lin = torch.optim.Adam(lin.parameters(), lr=1e-2)
for _ in range(200):
    opt_lin.zero_grad()
    nn.CrossEntropyLoss()(lin(h_tr), y_tr).backward(); opt_lin.step()
with torch.no_grad():
    acc_ssl = (lin(h_te).argmax(1)==y_te).float().mean().item()
print(f'SSL linear probe: {acc_ssl*100:.1f}%')   # 62.5%

34. Fill in: note for SimCLR pretraining + linear probe on…

Comparison

35. Without one step: The contrastive learning recipe

Constraint

36. The contrastive learning recipe

Pattern

1. Two augmented views of each image: random crop + flip + color-jitter + grayscale. Both views have the same label (unknown); different images are negatives.
2. Encode: h_i = f(x_i_aug1), h_j = f(x_i_aug2) — both through the same encoder f
3. Project: z_i = g(h_i) via a 2-layer MLP projection head. L2-normalize z before loss.
4. NT-Xent loss: stack [z1; z2] into (2B, D), compute sim = z @ zT / tau, mask diagonal, cross-entropy with labels [B..2B-1, 0..B-1]
5. Temperature tau: use 0.1 (SimCLR default). Too small → exploding gradients on hard negatives; too large → uniform distribution, no gradient signal.
6. Evaluate: freeze encoder, train nn.Linear(h_dim, n_classes) — this is the linear probe. Do NOT attach to z.
7. Mask R-CNN checklist: RoIAlign (not RoIPool) for mask branch; mask head output shape (N_proposals, K, 28, 28); use sigmoid per class, NOT softmax.

37. Where does each piece belong: Lesson 112: Instance Segmentation &…

Sorting

38. Rule out three: Check — InfoNCE temperature

Elimination

39. Check — InfoNCE temperature

Check

Work through the temperature effect before clicking.

40. Answer it before you see the options: Check — mask head output shape

Prediction

41. Check — mask head output shape

Check

Trace the shapes before clicking.

42. Rule out three: Check — linear probe vs fine-tuning

Elimination

43. Check — linear probe vs fine-tuning

Check

Apply the representation quality argument.

44. Your turn: SimCLR from scratch

Section

Project

45. Project: implement InfoNCE + SimCLR pipeline

Concept

Build a complete SimCLR training loop on synthetic 64-dim data. Three milestones: InfoNCE loss → full SimCLR pretrain → linear probe evaluation.

Build rules: L2-normalize embeddings before any cosine similarity; mask the diagonal (self-similarity) from the (2B, 2B) matrix; attach the linear probe to h = encoder(x), NOT to z = projector(h).

46. Break it if you can: Project: implement InfoNCE + SimCLR pipeline

Counterexample

47. Milestone 1 — NT-Xent loss implementation

Worked example

Your turn: implement nt_xent_loss. Key question: what are the labels in the cross-entropy call? For index i in [0..B-1], the positive is at index i+B in the 2B row; for index i in [B..2B-1], the positive is i-B.

Hint: torch.cat([torch.arange(B, 2*B), torch.arange(B)]) builds the label vector directly. After masking the diagonal with -1e9, nn.CrossEntropyLoss handles the rest.

import torch, torch.nn as nn

def nt_xent_loss(z1, z2, tau=0.5):
    B = z1.shape[0]
    z1 = nn.functional.normalize(z1, dim=-1)
    z2 = nn.functional.normalize(z2, dim=-1)
    z  = torch.cat([z1, z2], dim=0)      # (2B, D)
    sim = z @ z.T / tau                   # (2B, 2B)
    sim.masked_fill_(torch.eye(2*B, dtype=torch.bool), -1e9)
    labels = torch.cat([torch.arange(B, 2*B), torch.arange(B)])
    return nn.CrossEntropyLoss()(sim, labels)

torch.manual_seed(42)
B, D = 6, 8
z1 = torch.randn(B, D)
z2 = z1 + 0.05 * torch.randn_like(z1)    # tight positives
print(f'tau=0.07: {nt_xent_loss(z1, z2, 0.07):.4f}')  # 0.0009
print(f'tau=0.50: {nt_xent_loss(z1, z2, 0.50):.4f}')  # 0.8428
print(f'tau=1.00: {nt_xent_loss(z1, z2, 1.00):.4f}')  # 1.4959

48. Milestones 2 & 3 — pretrain + linear probe

Worked example

Your turn: add data augmentation (Gaussian noise views), run 60 epochs of SSL, then attach a frozen-encoder linear probe. Predict the final test accuracy before running — is 60 epochs enough to beat supervised?

Hint: augment each batch with x + 0.1*torch.randn_like(x) twice (different seeds). After pretraining, call encoder.eval() and extract h_tr = encoder(X_tr) with torch.no_grad() before training the linear head.

import torch, torch.nn as nn
# (TinyEncoder, ProjectionHead, nt_xent_loss defined in earlier milestones)
# --- Data ---
torch.manual_seed(42)
vec0 = torch.zeros(64); vec0[0] = 2.0
vec1 = torch.zeros(64); vec1[0] = -2.0
X = torch.cat([torch.randn(100,64)+vec0, torch.randn(100,64)+vec1])
y = torch.cat([torch.zeros(100,dtype=torch.long), torch.ones(100,dtype=torch.long)])
perm = torch.randperm(200, generator=torch.Generator().manual_seed(42))
X_tr, X_te = X[perm[:160]], X[perm[160:]]
y_tr, y_te = y[perm[:160]], y[perm[160:]]
# --- SSL pretrain ---
torch.manual_seed(42)
enc = TinyEncoder(64, 32, 16); prj = ProjectionHead(16, 16, 8)
opt = torch.optim.Adam(list(enc.parameters())+list(prj.parameters()), lr=1e-3)
for epoch in range(60):
    g1 = torch.Generator().manual_seed(epoch)
    g2 = torch.Generator().manual_seed(epoch+100)
    v1 = enc(X_tr + 0.1*torch.randn(X_tr.shape, generator=g1))
    v2 = enc(X_tr + 0.1*torch.randn(X_tr.shape, generator=g2))
    l = nt_xent_loss(prj(v1), prj(v2), tau=0.5)
    opt.zero_grad(); l.backward(); opt.step()
# --- Linear probe ---
with torch.no_grad():
    h_tr = enc(X_tr); h_te = enc(X_te)   # attach to h, NOT z
lin = nn.Linear(16, 2)
opt2 = torch.optim.Adam(lin.parameters(), lr=1e-2)
for _ in range(200):
    opt2.zero_grad()
    nn.CrossEntropyLoss()(lin(h_tr),y_tr).backward(); opt2.step()
with torch.no_grad():
    acc = (lin(h_te).argmax(1)==y_te).float().mean()
print(f'SSL linear probe: {acc*100:.1f}%')   # 62.5%

49. What each one costs: Milestones 2 & 3 — pretrain + linear probe

Trade off

50. The full program

Concept

import torch, torch.nn as nn

class TinyEncoder(nn.Module):
    def __init__(self, in_d=64, hid=32, out_d=16):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(in_d,hid), nn.ReLU(), nn.Linear(hid,out_d))
    def forward(self, x): return self.net(x)

class ProjectionHead(nn.Module):
    def __init__(self, in_d=16, hid=16, out_d=8):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(in_d,hid), nn.ReLU(), nn.Linear(hid,out_d))
    def forward(self, x): return self.net(x)

def nt_xent(z1, z2, tau=0.5):
    B = z1.shape[0]
    z1 = nn.functional.normalize(z1, dim=-1)
    z2 = nn.functional.normalize(z2, dim=-1)
    z  = torch.cat([z1, z2])
    sim = z @ z.T / tau
    sim.masked_fill_(torch.eye(2*B, dtype=torch.bool), -1e9)
    labels = torch.cat([torch.arange(B, 2*B), torch.arange(B)])
    return nn.CrossEntropyLoss()(sim, labels)

torch.manual_seed(42)
v0=torch.zeros(64); v0[0]=2.0; v1=torch.zeros(64); v1[0]=-2.0
X=torch.cat([torch.randn(100,64)+v0, torch.randn(100,64)+v1])
y=torch.cat([torch.zeros(100,dtype=torch.long), torch.ones(100,dtype=torch.long)])
perm=torch.randperm(200, generator=torch.Generator().manual_seed(42))
Xt,Xe,yt,ye=X[perm[:160]],X[perm[160:]],y[perm[:160]],y[perm[160:]]

torch.manual_seed(42)
enc=TinyEncoder(); prj=ProjectionHead()
opt=torch.optim.Adam(list(enc.parameters())+list(prj.parameters()), lr=1e-3)
for ep in range(60):
    g1=torch.Generator().manual_seed(ep)
    g2=torch.Generator().manual_seed(ep+100)
    l=nt_xent(prj(enc(Xt+0.1*torch.randn(Xt.shape,generator=g1))),
              prj(enc(Xt+0.1*torch.randn(Xt.shape,generator=g2))), tau=0.5)
    opt.zero_grad(); l.backward(); opt.step()
    if ep%10==0 or ep==59: print(f'ep {ep}: {l.item():.4f}')

with torch.no_grad(): ht=enc(Xt); he=enc(Xe)
lin=nn.Linear(16,2); opt2=torch.optim.Adam(lin.parameters(),lr=1e-2)
for _ in range(200):
    opt2.zero_grad(); nn.CrossEntropyLoss()(lin(ht),yt).backward(); opt2.step()
with torch.no_grad():
    print(f'Linear probe: {(lin(he).argmax(1)==ye).float().mean()*100:.1f}%')  # 62.5%

51. Fill in: rationale for The full program

Comparison

52. Show it off

Concept

Out loud, slides closed: (1) describe the full SimCLR pipeline from a raw batch of images to a scalar NT-Xent loss, naming every tensor shape; (2) explain why you attach the linear probe to h = f(x) and not to z = g(f(x)); (3) explain what happens to NT-Xent loss as tau → 0 and as tau → ∞.

Stretch (from the lesson plan): implement the InfoNCE mutual information lower bound — show that minimizing InfoNCE is equivalent to maximizing a lower bound on I(h; x). Also: run the linear probe for multiple label fractions (10%, 25%, 50%, 100%) and plot SSL vs supervised accuracy as a function of label count.

53. Connect it up: Lesson 112: Instance Segmentation & Contrastive Learning

Connect it up

54. What you can do now

Recap

• Trace Mask R-CNN's mask head: (N, 256, 14, 14) → conv stack → ConvTranspose2d → (N, K, 28, 28) — and explain why RoIAlign is required instead of RoIPool
• Distinguish semantic / instance / panoptic segmentation and state the PQ = SQ × RQ panoptic quality formula
• Implement NT-Xent / InfoNCE from scratch: L2-normalize → stack (2B, D) → cosine sim matrix / tau → mask diagonal → cross-entropy with symmetric labels
• Describe the SimCLR pipeline: augment → encoder f → projection head g → NT-Xent; discard g and probe on h after pretraining
• State the temperature effect: small tau sharpens softmax (good for easy pairs, explosive for hard); large tau loses gradient signal