File size: 4,250 Bytes
ecc201f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 | """
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")
|