Self-Supervised Vision

Module: Computer Vision | Difficulty: Advanced

Contrastive Learning

NT-Xent Loss (SimCLR)

MoCo Momentum Update

Masked Image Modeling (MAE)

BYOL (No Negative Pairs)

import torch
import torch.nn as nn

class SimCLR(nn.Module):
    def __init__(self, backbone, proj_dim=128):
        super().__init__()
        feat_dim = backbone.fc.in_features
        backbone.fc = nn.Identity()
        self.backbone = backbone
        self.projector = nn.Sequential(
            nn.Linear(feat_dim, feat_dim), nn.ReLU(True),
            nn.Linear(feat_dim, proj_dim),
        )
    
    def forward(self, x1, x2):
        h1 = self.projector(self.backbone(x1))
        h2 = self.projector(self.backbone(x2))
        return h1, h2

def nt_xent_loss(z1, z2, temperature=0.5):
    B = z1.size(0)
    z = torch.cat([z1, z2], dim=0)
    sim = nn.functional.cosine_similarity(z.unsqueeze(1), z.unsqueeze(0), dim=2) / temperature
    
    labels = torch.cat([torch.arange(B, 2*B), torch.arange(0, B)]).to(z.device)
    mask = torch.eye(2*B, dtype=torch.bool, device=z.device)
    sim.masked_fill_(mask, -1e9)
    
    return nn.functional.cross_entropy(sim, labels)

Key Takeaways

• Contrastive learning learns without labels by comparing views
• MAE masks random patches and learns to reconstruct
• Self-supervised pre-training achieves strong transfer performance