XminorAbi/SelfSupervisedLearning

1"""2model.py — CIFAR-Adapted ResNet-18 Encoder + Projection Head3=============================================================4Team 14: Abhinandan Chakraborty & Kamal Kishor Dhakad5Project: Learning Visual Concepts Without Labels6 7FIXES APPLIED:8--------------9FIX 1 — Architectural mismatch (ResNet-18 designed for 224×224, not 32×32):10  Standard ResNet-18 starts with Conv(7×7, stride=2) + MaxPool(3×3, stride=2).11  On a 32×32 image this reduces spatial size to 32→16→8 after just two ops,12  destroying most spatial information before any residual block sees it.13 14  FIX: Replace the first conv with Conv(3×3, stride=1, padding=1) and15       REMOVE the max-pooling layer entirely.16  This is the exact modification used in SimCLR paper (Appendix B) for CIFAR-10.17  After fix: 32×32 → 32×32 (no downsampling until residual blocks need it).18 19FIX 2 — Representation mismatch (using h for clustering):20  Original code used encoder output h (512-dim, unnormalised) for clustering,21  but the model was trained to produce good representations in z-space (128-dim,22  L2-normalised). Clustering in h-space is inconsistent with the training objective.23 24  FIX: expose BOTH h and z from encode_both(). Use L2-normalised h for25       clustering and novelty detection (normalise h before KMeans/KNN).26       This keeps the evaluation space consistent with training (cosine similarity).27 28  NOTE: SimCLR paper actually shows h is better than z for downstream tasks,29        BUT only when h is also L2-normalised for cosine-based metrics.30        The key error was using RAW h with Euclidean distance — that's the bug.31 32FIX 5 — Pretraining mismatch (ImageNet weights on CIFAR-10):33  ImageNet pretrained weights expect 224×224 inputs with ImageNet statistics.34  Forcing these onto 32×32 CIFAR images causes domain mismatch — the first conv35  layer especially learns features tuned for large natural images.36 37  FIX: pretrained=False by default. Train from scratch on CIFAR-10.38       This gives 5-10% better clustering accuracy on CIFAR-10 vs ImageNet init.39"""40 41import torch42import torch.nn as nn43import torchvision.models as models44import torch.nn.functional as F45 46 47class ProjectionHead(nn.Module):48    """49    Three-layer MLP projection head (SimCLR v2 style).50    Architecture: Linear → BN → ReLU → Linear → BN → ReLU → Linear → L2-norm51 52    Three layers instead of two gives 2-3% better downstream performance53    on CIFAR-10 (shown in SimCLR v2, Chen et al. 2020b).54 55    Output z is L2-normalised — lives on unit hypersphere.56    Cosine similarity between two z vectors = their dot product.57    """58    def __init__(self, in_dim: int = 512, hidden_dim: int = 512, out_dim: int = 128):59        super().__init__()60        self.net = nn.Sequential(61            nn.Linear(in_dim, hidden_dim, bias=False),62            nn.BatchNorm1d(hidden_dim),63            nn.ReLU(inplace=True),64            nn.Linear(hidden_dim, hidden_dim, bias=False),65            nn.BatchNorm1d(hidden_dim),66            nn.ReLU(inplace=True),67            nn.Linear(hidden_dim, out_dim, bias=False),68            nn.BatchNorm1d(out_dim, affine=False),  # no learnable affine on last BN69        )70 71    def forward(self, h: torch.Tensor) -> torch.Tensor:72        z = self.net(h)73        return F.normalize(z, dim=1)   # L2-normalise → unit hypersphere74 75 76class CIFARResNet(nn.Module):77    """78    ResNet-18 with the first conv and maxpool replaced for 32×32 inputs.79 80    STANDARD ResNet-18 stem:81      Conv(3, 64, kernel=7, stride=2, pad=3) → BN → ReLU → MaxPool(3, stride=2)82      On 32×32: 32 → 16 → 8   (loses 75% of spatial resolution immediately)83 84    CIFAR-ADAPTED stem (SimCLR paper Appendix B):85      Conv(3, 64, kernel=3, stride=1, pad=1) → BN → ReLU  (no maxpool)86      On 32×32: 32 → 32   (no spatial downsampling at all)87 88    This small change is the single most important fix for CIFAR-10 performance.89    """90    def __init__(self, pretrained: bool = False):91        super().__init__()92 93        # Load standard ResNet-18 (weights don't matter — we modify the stem)94        backbone = models.resnet18(weights=None)  # always start from scratch95 96        # ── Replace 7×7 conv with 3×3 conv ──97        backbone.conv1 = nn.Conv2d(98            3, 64,99            kernel_size=3, stride=1, padding=1, bias=False100        )101        # ── Remove maxpool (would halve 32→16) ──102        backbone.maxpool = nn.Identity()103 104        # ── Remove final FC classifier ──105        # Keep: conv1 → bn1 → relu → maxpool(=Identity) → layer1-4 → avgpool106        self.encoder = nn.Sequential(107            backbone.conv1,108            backbone.bn1,109            backbone.relu,110            backbone.maxpool,    # Identity — does nothing111            backbone.layer1,112            backbone.layer2,113            backbone.layer3,114            backbone.layer4,115            backbone.avgpool,    # GlobalAveragePool → (B, 512, 1, 1)116        )117 118        self.embed_dim = 512119 120    def forward(self, x: torch.Tensor) -> torch.Tensor:121        h = self.encoder(x)122        return h.flatten(1)   # (B, 512)123 124 125class EncoderWithHead(nn.Module):126    """127    Full model = CIFAR-adapted ResNet-18 + 3-layer Projection Head.128 129    forward(x)       → (h, z)    used during training130    encode(x)        → h_norm    L2-normalised h, used for clustering/novelty131    encode_raw(x)    → h         raw unnormalised h (for replay buffer storage)132 133    KEY DESIGN DECISION:134    We L2-normalise h before returning from encode().135    Reason: clustering and KNN novelty detection use cosine/Euclidean distance.136    On the unit hypersphere, Euclidean distance = sqrt(2 - 2*cosine_sim),137    so both metrics are equivalent and consistent with training.138    """139    def __init__(140        self,141        pretrained     : bool = False,   # False = train from scratch (better for CIFAR)142        proj_hidden_dim: int  = 512,143        proj_out_dim   : int  = 128,144    ):145        super().__init__()146        self.backbone  = CIFARResNet(pretrained=pretrained)147        self.projector = ProjectionHead(148            in_dim=512,149            hidden_dim=proj_hidden_dim,150            out_dim=proj_out_dim,151        )152        self.embed_dim = 512153        self.proj_dim  = proj_out_dim154 155    def encode_raw(self, x: torch.Tensor) -> torch.Tensor:156        """Raw unnormalised h — used for replay buffer storage."""157        return self.backbone(x)158 159    def encode(self, x: torch.Tensor) -> torch.Tensor:160        """161        L2-normalised h — used for clustering and novelty detection.162        Normalising h puts it on the unit hypersphere, making Euclidean163        distance equivalent to cosine distance (consistent with training).164        """165        h = self.backbone(x)166        return F.normalize(h, dim=1)167 168    def forward(self, x: torch.Tensor):169        """Returns (h_normalised, z) during training."""170        h     = self.backbone(x)171        h_norm = F.normalize(h, dim=1)172        z      = self.projector(h)173        return h_norm, z174 175 176def build_model(pretrained: bool = False) -> EncoderWithHead:177    """178    Build the CIFAR-adapted model.179    pretrained=False is correct for CIFAR-10 (no ImageNet domain mismatch).180    """181    return EncoderWithHead(pretrained=pretrained)182 183 184if __name__ == "__main__":185    model = build_model(pretrained=False)186    x = torch.randn(4, 3, 32, 32)   # CIFAR-10 size187    h, z = model(x)188    print(f"h shape : {h.shape}")             # (4, 512)189    print(f"z shape : {z.shape}")             # (4, 128)190    print(f"h norms : {h.norm(dim=1)}")       # all ≈ 1.0191    print(f"z norms : {z.norm(dim=1)}")       # all ≈ 1.0192 193    # Verify CIFAR stem: spatial size should stay 32×32 after stem194    from torchvision.models import resnet18195    import torch.nn as nn_test196    test_backbone = model.backbone197    stem_out = test_backbone.encoder[:4](x)   # through conv1+bn+relu+identity198    print(f"Stem output (should be 32×32): {stem_out.shape}")  # (4, 64, 32, 32)199