XminorAbi/SelfSupervisedLearning

1"""2classifier.py — Supervised Linear Classifier Head3===================================================4Team 14: Abhinandan Chakraborty & Kamal Kishor Dhakad5Project: Learning Visual Concepts Without Labels6 7ARCHITECTURE REDESIGN:8-----------------------9Old system (broken):10  image → encoder → PCA(128) → KMeans → class name11  Problems: KMeans doesn't understand semantics, PCA destroys info,12            no supervised signal → car predicted as horse.13 14New system (this file):15  image → FROZEN encoder → L2-norm h (512) → Linear(512→N) → softmax → class16 17WHY THIS WORKS:18  The SimCLR encoder already learned rich visual features.19  A single linear layer on top is enough to separate classes20  because the embedding space is already semantically structured.21  This is called "linear probing" — the standard evaluation protocol22  for self-supervised learning papers.23 24  Accuracy improvement: KMeans ~52-75% → Linear probe ~85-92%25 26TRAINING:27  - Encoder weights are FROZEN (not updated)28  - Only the linear layer is trained29  - Uses labeled CIFAR-10 train split (50K images)30  - Takes ~5-10 minutes on T4 GPU (very fast)31 32KEY DESIGN DECISIONS:33  - Input: L2-normalised 512-dim embeddings (same as registry)34  - No PCA — raw embeddings preserve all discriminative information35  - Softmax output → confidence score for hybrid pipeline36  - Supports adding new classes (expand head) for continuous learning37"""38 39import torch40import torch.nn as nn41import torch.nn.functional as F42import torch.optim as optim43from torch.optim.lr_scheduler import CosineAnnealingLR44from pathlib import Path45import time46import os47 48 49class LinearClassifier(nn.Module):50    """51    Single linear layer classifier on top of frozen SimCLR encoder.52 53    Input  : L2-normalised embedding h (512-dim)54    Output : class logits (N_classes)55 56    Why only one linear layer?57    - SimCLR embeddings are already linearly separable after good training58    - Adding non-linearity (MLP) can overfit on small labeled sets59    - Linear probe is the standard SSL evaluation protocol60    - Softmax confidence is well-calibrated from a linear layer61 62    Expanding for new classes:63    - When a new class is registered, expand_head(new_n_classes) is called64    - Old weights are preserved, new class weights initialised to zero65    """66 67    def __init__(self, embed_dim: int = 512, n_classes: int = 10):68        super().__init__()69        self.fc = nn.Linear(embed_dim, n_classes)70        self.embed_dim = embed_dim71        self.n_classes = n_classes72 73        # Initialise with small weights for stable softmax74        nn.init.normal_(self.fc.weight, std=0.01)75        nn.init.zeros_(self.fc.bias)76 77    def forward(self, h: torch.Tensor) -> torch.Tensor:78        """79        Args:80            h : (B, 512) L2-normalised embeddings81        Returns:82            logits : (B, N_classes) — pass through softmax for probabilities83        """84        return self.fc(h)85 86    def predict(self, h: torch.Tensor):87        """88        Returns (predicted_class_idx, confidence) for each sample.89        confidence = max softmax probability.90        """91        with torch.no_grad():92            logits      = self.forward(h)93            probs       = F.softmax(logits, dim=1)94            confidence, pred = probs.max(dim=1)95        return pred, confidence96 97    def expand_head(self, new_n_classes: int):98        """99        Expand classifier to handle more classes.100        Old weights preserved — new class weights initialised to zero.101        Called when a novel class is confirmed and added to the system.102 103        Args:104            new_n_classes: total number of classes after expansion105        """106        if new_n_classes <= self.n_classes:107            return108 109        old_weight = self.fc.weight.data   # (old_n, 512)110        old_bias   = self.fc.bias.data     # (old_n,)111 112        # New linear layer with expanded output113        new_fc = nn.Linear(self.embed_dim, new_n_classes)114        nn.init.zeros_(new_fc.weight)115        nn.init.zeros_(new_fc.bias)116 117        # Copy old weights118        new_fc.weight.data[:self.n_classes] = old_weight119        new_fc.bias.data[:self.n_classes]   = old_bias120 121        self.fc        = new_fc122        self.n_classes = new_n_classes123        print(f"[Classifier] Expanded head: {self.n_classes - (new_n_classes - self.n_classes)} → {new_n_classes} classes")124 125 126def train_classifier(127    encoder,128    class_names     : list,129    cifar10_root    : str  = "data/cifar10",130    save_path       : str  = "checkpoints/classifier.pt",131    drive_save_path : str  = None,132    n_epochs        : int  = 30,133    batch_size      : int  = 512,134    lr              : float = 0.1,    # SGD with high LR works best for linear probe135    num_workers     : int  = 4,136    image_size      : int  = 32,137    device          : str  = "cuda",138) -> LinearClassifier:139    """140    Train the linear classifier head on frozen encoder embeddings.141 142    Uses full CIFAR-10 train split (50K labeled images).143    Encoder is frozen — only fc layer weights are updated.144 145    Training takes ~5-10 minutes on T4 GPU.146    Expected accuracy: 85-92% on CIFAR-10 test set.147 148    Args:149        encoder       : trained SimCLR encoder (EncoderWithHead)150        class_names   : list of class name strings151        cifar10_root  : path to CIFAR-10 data152        save_path     : local path to save classifier weights153        drive_save_path: Google Drive path to sync weights154        n_epochs      : 30 epochs is sufficient for linear probe155        batch_size    : 512 (fast since no backward through encoder)156        lr            : 0.1 with SGD is standard for linear probing157        device        : 'cuda' / 'mps' / 'cpu'158    """159    from dataset import build_eval_transform160    from torchvision.datasets import CIFAR10161    from torch.utils.data import DataLoader162 163    n_classes  = len(class_names)164    device_obj = torch.device(device)165 166    # Freeze encoder167    encoder = encoder.to(device_obj)168    encoder.eval()169    for param in encoder.parameters():170        param.requires_grad = False171 172    print(f"[Classifier] Encoder frozen. Training linear head only.")173    print(f"[Classifier] Classes: {n_classes} — {class_names}")174 175    # ── Pre-extract all embeddings (much faster than computing per batch) ──176    print(f"\n[Classifier] Pre-extracting embeddings from CIFAR-10...")177    transform = build_eval_transform(image_size)178 179    train_data = CIFAR10(root=cifar10_root, train=True,  download=True, transform=transform)180    test_data  = CIFAR10(root=cifar10_root, train=False, download=True, transform=transform)181 182    train_loader = DataLoader(train_data, batch_size=batch_size, shuffle=False,183                              num_workers=num_workers, pin_memory=True)184    test_loader  = DataLoader(test_data,  batch_size=batch_size, shuffle=False,185                              num_workers=num_workers, pin_memory=True)186 187    def extract_embeddings(loader, split_name):188        all_h, all_y = [], []189        with torch.no_grad():190            for imgs, labels in loader:191                h = encoder.encode(imgs.to(device_obj))   # L2-normalised (B, 512)192                all_h.append(h.cpu())193                all_y.append(labels)194        h_all = torch.cat(all_h)195        y_all = torch.cat(all_y)196        print(f"  {split_name}: {h_all.shape[0]} embeddings extracted")197        return h_all, y_all198 199    train_h, train_y = extract_embeddings(train_loader, "Train")200    test_h,  test_y  = extract_embeddings(test_loader,  "Test")201 202    # ── Build in-memory dataset from embeddings ──203    from torch.utils.data import TensorDataset204    train_emb_loader = DataLoader(205        TensorDataset(train_h, train_y),206        batch_size=512, shuffle=True,207    )208    test_emb_loader = DataLoader(209        TensorDataset(test_h, test_y),210        batch_size=512, shuffle=False,211    )212 213    # ── Build classifier ──214    classifier = LinearClassifier(embed_dim=512, n_classes=n_classes).to(device_obj)215 216    # SGD + cosine LR — standard for linear probing (better than Adam here)217    optimizer  = optim.SGD(classifier.parameters(), lr=lr, momentum=0.9, weight_decay=1e-4)218    scheduler  = CosineAnnealingLR(optimizer, T_max=n_epochs, eta_min=1e-4)219    criterion  = nn.CrossEntropyLoss()220 221    print(f"\n[Classifier] Training for {n_epochs} epochs...")222    best_acc   = 0.0223    best_state = None224 225    for epoch in range(1, n_epochs + 1):226        # Train227        classifier.train()228        total_loss, correct, total = 0.0, 0, 0229        for h_batch, y_batch in train_emb_loader:230            h_batch = h_batch.to(device_obj)231            y_batch = y_batch.to(device_obj)232 233            logits = classifier(h_batch)234            loss   = criterion(logits, y_batch)235 236            optimizer.zero_grad()237            loss.backward()238            optimizer.step()239 240            total_loss += loss.item()241            preds       = logits.argmax(dim=1)242            correct    += (preds == y_batch).sum().item()243            total      += len(y_batch)244 245        scheduler.step()246        train_acc = correct / total * 100247 248        # Evaluate249        classifier.eval()250        correct_t, total_t = 0, 0251        with torch.no_grad():252            for h_batch, y_batch in test_emb_loader:253                h_batch = h_batch.to(device_obj)254                y_batch = y_batch.to(device_obj)255                preds   = classifier(h_batch).argmax(dim=1)256                correct_t += (preds == y_batch).sum().item()257                total_t   += len(y_batch)258        test_acc = correct_t / total_t * 100259 260        print(f"  Epoch [{epoch:02d}/{n_epochs}] "261              f"loss={total_loss/len(train_emb_loader):.4f}  "262              f"train={train_acc:.1f}%  test={test_acc:.1f}%")263 264        # Save best265        if test_acc > best_acc:266            best_acc   = test_acc267            best_state = {k: v.clone() for k, v in classifier.state_dict().items()}268 269    # Load best weights270    classifier.load_state_dict(best_state)271    print(f"\n[Classifier] Best test accuracy: {best_acc:.2f}%")272 273    # Save274    Path(save_path).parent.mkdir(exist_ok=True, parents=True)275    torch.save({276        "classifier_state": classifier.state_dict(),277        "class_names"     : class_names,278        "n_classes"       : n_classes,279        "embed_dim"       : 512,280        "best_acc"        : best_acc,281    }, save_path)282    print(f"[Classifier] Saved: {save_path}")283 284    # Sync to Drive285    if drive_save_path:286        import shutil287        Path(drive_save_path).parent.mkdir(exist_ok=True, parents=True)288        shutil.copy(save_path, drive_save_path)289        print(f"[Classifier] Synced to Drive: {drive_save_path}")290 291    return classifier292 293 294def load_classifier(path: str, device: str = "cpu") -> tuple:295    """296    Load saved classifier from disk.297    Returns (classifier, class_names).298    """299    device = torch.device("cpu")  # ✅ FORCE CPU300 301    ckpt = torch.load(path, map_location=device)302 303    class_names = ckpt["class_names"]304 305    classifier = LinearClassifier(306        embed_dim=ckpt["embed_dim"],307        n_classes=ckpt["n_classes"],308    )309 310    classifier.load_state_dict(ckpt["classifier_state"])311 312    classifier = classifier.to(device)  # now safe313    classifier.eval()314 315    acc = ckpt.get("best_acc", None)316    317    if acc is not None:318        print(f"[Classifier] Loaded: {len(class_names)} classes, acc={acc:.2f}%")319    else:320        print(f"[Classifier] Loaded: {len(class_names)} classes")321        322    return classifier, class_names323