Upload student_unet.py with huggingface_hub

student_unet.py

@@ -0,0 +1,114 @@
+"""
+student_unet.py
+MobileNetV2 + U-Net skip connections student model for DLoc.
+
+Architecture (Professor Phase 2, Step 5 Option 1):
+  Encoder: MobileNetV2 backbone split into 4 stages for skip connections
+    - inc:  ConvBNReLU(4->32, stride=2)
+    - enc1: features[1:4]   -> 24 channels  (Skip 1)
+    - enc2: features[4:7]   -> 32 channels  (Skip 2)
+    - enc3: features[7:14]  -> 96 channels  (Skip 3)
+    - enc4: features[14:18] -> 320 channels (Bottleneck)
+  Decoder: Upsample + Concat skip + Conv
+    - dec1: 320+96=416 -> 128
+    - dec2: 128+32=160 -> 64
+    - dec3: 64+24=88   -> 32
+    - final: 32 -> 1 + Sigmoid
+  Output: [B, 1, 161, 361]
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision.models import mobilenet_v2, MobileNet_V2_Weights
+
+
+class MobileNetUNetStudent(nn.Module):
+    OUTPUT_SIZE = (161, 361)
+
+    def __init__(self, input_nc=4):
+        super().__init__()
+
+        mobilenet = mobilenet_v2(weights=MobileNet_V2_Weights.DEFAULT)
+
+        # Modify first conv for input_nc channels
+        orig_conv = mobilenet.features[0][0]
+        self.inc = nn.Sequential(
+            nn.Conv2d(input_nc, orig_conv.out_channels,
+                      kernel_size=orig_conv.kernel_size,
+                      stride=orig_conv.stride,
+                      padding=orig_conv.padding, bias=False),
+            mobilenet.features[0][1],  # BatchNorm
+            mobilenet.features[0][2],  # ReLU6
+        )
+
+        # Encoder stages (split to capture skip connections)
+        self.enc1 = mobilenet.features[1:4]    # -> 24 ch
+        self.enc2 = mobilenet.features[4:7]    # -> 32 ch
+        self.enc3 = mobilenet.features[7:14]   # -> 96 ch
+        self.enc4 = mobilenet.features[14:-1]  # -> 320 ch (drop final 1x1 conv)
+
+        # Decoder with skip connections
+        self.up1 = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False)
+        self.dec1 = nn.Sequential(
+            nn.Conv2d(320 + 96, 128, 3, padding=1, bias=False),
+            nn.BatchNorm2d(128),
+            nn.ReLU(inplace=True),
+        )
+
+        self.up2 = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False)
+        self.dec2 = nn.Sequential(
+            nn.Conv2d(128 + 32, 64, 3, padding=1, bias=False),
+            nn.BatchNorm2d(64),
+            nn.ReLU(inplace=True),
+        )
+
+        self.up3 = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False)
+        self.dec3 = nn.Sequential(
+            nn.Conv2d(64 + 24, 32, 3, padding=1, bias=False),
+            nn.BatchNorm2d(32),
+            nn.ReLU(inplace=True),
+        )
+
+        self.final_conv = nn.Sequential(
+            nn.Conv2d(32, 1, kernel_size=1),
+            nn.Sigmoid(),
+        )
+
+    def forward(self, x):
+        # Encoder
+        x0 = self.inc(x)
+        s1 = self.enc1(x0)       # skip 1: 24 ch
+        s2 = self.enc2(s1)       # skip 2: 32 ch
+        s3 = self.enc3(s2)       # skip 3: 96 ch
+        bottle = self.enc4(s3)   # bottleneck: 320 ch
+
+        # Decoder with skip concatenation
+        d1 = self.up1(bottle)
+        s3r = F.interpolate(s3, size=d1.shape[2:], mode='bilinear', align_corners=False)
+        d1 = self.dec1(torch.cat([d1, s3r], dim=1))
+
+        d2 = self.up2(d1)
+        s2r = F.interpolate(s2, size=d2.shape[2:], mode='bilinear', align_corners=False)
+        d2 = self.dec2(torch.cat([d2, s2r], dim=1))
+
+        d3 = self.up3(d2)
+        s1r = F.interpolate(s1, size=d3.shape[2:], mode='bilinear', align_corners=False)
+        d3 = self.dec3(torch.cat([d3, s1r], dim=1))
+
+        out = self.final_conv(d3)
+        if out.shape[2:] != self.OUTPUT_SIZE:
+            out = F.interpolate(out, size=self.OUTPUT_SIZE, mode='bilinear', align_corners=False)
+        return out
+
+
+if __name__ == "__main__":
+    model = MobileNetUNetStudent(input_nc=4)
+    n = sum(p.numel() for p in model.parameters())
+    print(f"Parameters: {n:,} ({n/1e6:.2f}M)")
+    x = torch.randn(2, 4, 161, 361)
+    with torch.no_grad():
+        y = model(x)
+    print(f"Input:  {x.shape}")
+    print(f"Output: {y.shape}")
+    assert y.shape == (2, 1, 161, 361), f"Wrong shape: {y.shape}"
+    print("OK")