Upload train_mse_distill.py

train_mse_distill.py

@@ -0,0 +1,395 @@
+"""
+Fixed distillation training for microbubble segmentation.
+
+KEY FIX: Uses MSE distillation on teacher's RAW cell_prob LOGITS instead of
+BCE on binary masks. This solves the class imbalance problem (foreground is
+only ~0.2% of pixels) that causes BCE training to predict all-background.
+
+Usage:
+    # 1. Generate pseudo-labels (already done for callumtilbury/microbubble-images)
+    python generate_pseudolabels.py \
+        --image_dir ./images \
+        --model_path ./teacher_model/bubble_finetuned \
+        --output_dir ./pseudolabels
+
+    # 2. Train student with MSE distillation
+    python train_mse_distill.py \
+        --image_dir ./images \
+        --label_dir ./pseudolabels \
+        --output_dir ./checkpoints_v3 \
+        --epochs 400 \
+        --batch_size 4 \
+        --crop_size 256
+
+    # 3. Push to hub
+    python train_mse_distill.py --push_to_hub callumtilbury/bubble-distill-v3
+"""
+
+import argparse
+import json
+import time
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader, random_split
+from skimage import io as skio
+from huggingface_hub import HfApi, create_repo
+
+
+# ============================================================================
+# TinyBubbleNet (depthwise-separable U-Net, ~389K params)
+# ============================================================================
+
+class DepthwiseSeparableConv(nn.Module):
+    def __init__(self, in_ch, out_ch, kernel_size=3, padding=1, stride=1):
+        super().__init__()
+        self.depthwise = nn.Conv2d(in_ch, in_ch, kernel_size, stride=stride, padding=padding, groups=in_ch, bias=False)
+        self.pointwise = nn.Conv2d(in_ch, out_ch, 1, bias=False)
+        self.bn = nn.BatchNorm2d(out_ch)
+        self.relu = nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        return self.relu(self.bn(self.pointwise(self.depthwise(x))))
+
+
+class ConvBlock(nn.Module):
+    def __init__(self, in_ch, out_ch, use_depthwise=True):
+        super().__init__()
+        if use_depthwise:
+            self.block = nn.Sequential(
+                DepthwiseSeparableConv(in_ch, out_ch),
+                DepthwiseSeparableConv(out_ch, out_ch),
+            )
+        else:
+            self.block = nn.Sequential(
+                nn.Conv2d(in_ch, out_ch, 3, padding=1, bias=False),
+                nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True),
+                nn.Conv2d(out_ch, out_ch, 3, padding=1, bias=False),
+                nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True),
+            )
+
+    def forward(self, x):
+        return self.block(x)
+
+
+class TinyBubbleNet(nn.Module):
+    def __init__(self, in_channels=1, base_ch=16, out_channels=4, use_depthwise=True):
+        super().__init__()
+        self.enc1 = ConvBlock(in_channels, base_ch, use_depthwise=False)
+        self.enc2 = ConvBlock(base_ch, base_ch * 2, use_depthwise)
+        self.enc3 = ConvBlock(base_ch * 2, base_ch * 4, use_depthwise)
+        self.enc4 = ConvBlock(base_ch * 4, base_ch * 8, use_depthwise)
+        self.pool = nn.MaxPool2d(2)
+        self.bottleneck = ConvBlock(base_ch * 8, base_ch * 16, use_depthwise)
+        self.up4 = nn.ConvTranspose2d(base_ch * 16, base_ch * 8, kernel_size=2, stride=2)
+        self.dec4 = ConvBlock(base_ch * 16, base_ch * 8, use_depthwise)
+        self.up3 = nn.ConvTranspose2d(base_ch * 8, base_ch * 4, kernel_size=2, stride=2)
+        self.dec3 = ConvBlock(base_ch * 8, base_ch * 4, use_depthwise)
+        self.up2 = nn.ConvTranspose2d(base_ch * 4, base_ch * 2, kernel_size=2, stride=2)
+        self.dec2 = ConvBlock(base_ch * 4, base_ch * 2, use_depthwise)
+        self.up1 = nn.ConvTranspose2d(base_ch * 2, base_ch, kernel_size=2, stride=2)
+        self.dec1 = ConvBlock(base_ch * 2, base_ch, use_depthwise)
+        self.flow_head = nn.Conv2d(base_ch, 2, 1)
+        self.prob_head = nn.Conv2d(base_ch, 1, 1)
+        self.dist_head = nn.Conv2d(base_ch, 1, 1)
+
+    def forward(self, x):
+        e1 = self.enc1(x)
+        e2 = self.enc2(self.pool(e1))
+        e3 = self.enc3(self.pool(e2))
+        e4 = self.enc4(self.pool(e3))
+        b = self.bottleneck(self.pool(e4))
+        d4 = self.dec4(torch.cat([self.up4(b), e4], dim=1))
+        d3 = self.dec3(torch.cat([self.up3(d4), e3], dim=1))
+        d2 = self.dec2(torch.cat([self.up2(d3), e2], dim=1))
+        d1 = self.dec1(torch.cat([self.up1(d2), e1], dim=1))
+        return torch.cat([self.flow_head(d1), self.prob_head(d1), self.dist_head(d1)], dim=1)
+
+    def predict(self, x):
+        with torch.no_grad():
+            out = self.forward(x)
+        return {
+            "dY": out[:, 0], "dX": out[:, 1],
+            "cell_prob": torch.sigmoid(out[:, 2]),
+            "dist_transform": torch.relu(out[:, 3]),
+        }
+
+
+# ============================================================================
+# Dataset — uses teacher's RAW cell_prob LOGITS (not binary mask)
+# ============================================================================
+
+class RawPseudoLabelDataset(torch.utils.data.Dataset):
+    """
+    Loads image + pseudo-label pairs where cell_prob is the RAW LOGITS
+    from the teacher (range roughly -9 to +5), not a binary mask.
+    
+    This is the KEY FIX: MSE distillation on logits gives gradients on
+    ALL pixels, solving the class imbalance (foreground ~0.2%) that
+    breaks BCE-based training.
+    """
+    def __init__(self, image_dir, label_dir, crop_size=None, augment=True, normalize=True):
+        self.image_dir = Path(image_dir)
+        self.label_dir = Path(label_dir)
+        self.crop_size = crop_size
+        self.augment = augment
+        self.normalize = normalize
+        self.samples = []
+        for f in self.label_dir.glob("*_flows.npy"):
+            stem = f.stem.replace("_flows", "")
+            img_path = self.image_dir / f"{stem}.png"
+            if not img_path.exists():
+                for ext in [".jpg", ".jpeg", ".tif", ".tiff", ".bmp"]:
+                    alt = self.image_dir / f"{stem}{ext}"
+                    if alt.exists():
+                        img_path = alt
+                        break
+            if img_path.exists():
+                self.samples.append({
+                    "image_path": img_path,
+                    "flows_path": self.label_dir / f"{stem}_flows.npy",
+                    "dist_path": self.label_dir / f"{stem}_dist.npy",
+                })
+        print(f"Found {len(self.samples)} samples")
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, idx):
+        s = self.samples[idx]
+        img = skio.imread(str(s["image_path"]))
+        if img.ndim == 3 and img.shape[2] >= 3:
+            img = np.mean(img[:, :, :3], axis=2)
+        img = img.astype(np.float32)
+        flows = np.load(s["flows_path"]).astype(np.float32)
+        dist = np.load(s["dist_path"]).astype(np.float32)
+        if self.normalize:
+            p1, p99 = np.percentile(img, [1, 99])
+            img = np.clip((img - p1) / (p99 - p1 + 1e-8), 0, 1) if p99 > p1 else img / (img.max() + 1e-8)
+        flow_dY, flow_dX, cell_prob_logits = flows[0], flows[1], flows[2]
+        dist_max = dist.max()
+        dist_norm = dist / (dist_max + 1e-8) if dist_max > 0 else dist
+        target = np.stack([flow_dY, flow_dX, cell_prob_logits, dist_norm], axis=0)
+        if self.augment:
+            img, target = self._augment(img, target)
+        if self.crop_size is not None:
+            img, target = self._random_crop(img, target, self.crop_size)
+        return torch.from_numpy(img[np.newaxis]).float(), torch.from_numpy(target).float()
+
+    def _augment(self, img, target):
+        if np.random.random() < 0.5:
+            img = np.flip(img, axis=1).copy()
+            target = np.flip(target, axis=2).copy()
+            target[1] = -target[1]
+        if np.random.random() < 0.5:
+            img = np.flip(img, axis=0).copy()
+            target = np.flip(target, axis=1).copy()
+            target[0] = -target[0]
+        k = np.random.randint(4)
+        if k > 0:
+            img = np.rot90(img, k).copy()
+            target = np.rot90(target, k, axes=(1, 2)).copy()
+            if k == 1:
+                target[0], target[1] = target[1].copy(), -target[0].copy()
+            elif k == 2:
+                target[0], target[1] = -target[0], -target[1]
+            elif k == 3:
+                target[0], target[1] = -target[1].copy(), target[0].copy()
+        if np.random.random() < 0.5:
+            img = np.clip(img + np.random.uniform(-0.1, 0.1), 0, 1)
+        if np.random.random() < 0.5:
+            m = img.mean()
+            img = np.clip((img - m) * np.random.uniform(0.8, 1.2) + m, 0, 1)
+        if np.random.random() < 0.3:
+            img = np.clip(img + np.random.normal(0, 0.02, img.shape).astype(np.float32), 0, 1)
+        return img, target
+
+    def _random_crop(self, img, target, crop_size):
+        h, w = img.shape
+        ch, cw = crop_size
+        if h <= ch or w <= cw:
+            ph, pw = max(ch - h, 0), max(cw - w, 0)
+            img = np.pad(img, ((0, ph), (0, pw)), mode="reflect")
+            target = np.pad(target, ((0, 0), (0, ph), (0, pw)), mode="reflect")
+            h, w = img.shape
+        y = np.random.randint(0, h - ch + 1)
+        x = np.random.randint(0, w - cw + 1)
+        return img[y:y+ch, x:x+cw], target[:, y:y+ch, x:x+cw]
+
+
+# ============================================================================
+# Loss — pure MSE distillation on all 4 channels (the fix!)
+# ============================================================================
+
+class MSEDistillationLoss(nn.Module):
+    """
+    MSE distillation on teacher logits for ALL channels.
+    
+    Why this works:
+    - BCE on binary masks fails because foreground is only ~0.2% of pixels.
+      The model can get 99.8% accuracy by predicting ALL background.
+    - MSE on teacher logits gives gradients on ALL pixels. Even background
+      pixels have informative negative logit values (~ -6) that the student
+      must learn to reproduce.
+    """
+    def __init__(self, flow_weight=1.0, prob_weight=1.0, dist_weight=1.0):
+        super().__init__()
+        self.flow_weight = flow_weight
+        self.prob_weight = prob_weight
+        self.dist_weight = dist_weight
+
+    def forward(self, pred, target):
+        flow = F.mse_loss(pred[:, :2], target[:, :2])
+        prob = F.mse_loss(pred[:, 2:3], target[:, 2:3])
+        dist = F.mse_loss(torch.relu(pred[:, 3:4]), target[:, 3:4])
+        total = self.flow_weight * flow + self.prob_weight * prob + self.dist_weight * dist
+        return total, {"total": total.item(), "flow": flow.item(), "prob": prob.item(), "dist": dist.item()}
+
+
+# ============================================================================
+# Training loop
+# ============================================================================
+
+def train(args):
+    device = torch.device("cuda" if torch.cuda.is_available() and not args.no_gpu else "cpu")
+    print(f"Device: {device}")
+
+    dataset = RawPseudoLabelDataset(
+        image_dir=args.image_dir, label_dir=args.label_dir,
+        crop_size=(args.crop_size, args.crop_size) if args.crop_size else None,
+        augment=True, normalize=True,
+    )
+
+    n_val = max(1, int(len(dataset) * args.val_split))
+    n_train = len(dataset) - n_val
+    train_set, val_set = random_split(
+        dataset, [n_train, n_val], generator=torch.Generator().manual_seed(42))
+
+    train_loader = DataLoader(train_set, batch_size=args.batch_size, shuffle=True,
+                              num_workers=0, pin_memory=True, drop_last=True)
+    val_loader = DataLoader(val_set, batch_size=args.batch_size, shuffle=False,
+                            num_workers=0, pin_memory=True)
+
+    print(f"Train: {n_train}, Val: {n_val}")
+
+    model = TinyBubbleNet(in_channels=1, base_ch=args.base_ch, out_channels=4,
+                          use_depthwise=args.use_depthwise).to(device)
+    n_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    print(f"Params: {n_params:,}")
+
+    criterion = MSEDistillationLoss(flow_weight=args.flow_weight,
+                                    prob_weight=args.prob_weight,
+                                    dist_weight=args.dist_weight)
+    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
+    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
+        optimizer, T_max=args.epochs, eta_min=args.lr * 0.01)
+
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    config = vars(args)
+    config["n_params"] = n_params
+    with open(output_dir / "config.json", "w") as f:
+        json.dump(config, f, indent=2)
+
+    best_val_loss = float("inf")
+    history = []
+
+    for epoch in range(1, args.epochs + 1):
+        t0 = time.time()
+        model.train()
+        train_losses = {"total": 0.0, "flow": 0.0, "prob": 0.0, "dist": 0.0}
+        for images, targets in train_loader:
+            images, targets = images.to(device), targets.to(device)
+            pred = model(images)
+            loss, ld = criterion(pred, targets)
+            optimizer.zero_grad()
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
+            optimizer.step()
+            for k in train_losses:
+                train_losses[k] += ld[k]
+        for k in train_losses:
+            train_losses[k] /= len(train_loader)
+
+        model.eval()
+        val_losses = {"total": 0.0, "flow": 0.0, "prob": 0.0, "dist": 0.0}
+        with torch.no_grad():
+            for images, targets in val_loader:
+                images, targets = images.to(device), targets.to(device)
+                pred = model(images)
+                _, ld = criterion(pred, targets)
+                for k in val_losses:
+                    val_losses[k] += ld[k]
+        for k in val_losses:
+            val_losses[k] /= max(len(val_loader), 1)
+
+        scheduler.step()
+        elapsed = time.time() - t0
+        lr = optimizer.param_groups[0]["lr"]
+
+        print(f"Epoch {epoch:04d}/{args.epochs} | "
+              f"Train: {train_losses['total']:.4f} (f={train_losses['flow']:.4f}, "
+              f"p={train_losses['prob']:.4f}, d={train_losses['dist']:.4f}) | "
+              f"Val: {val_losses['total']:.4f} | LR: {lr:.2e} | {elapsed:.1f}s")
+
+        if val_losses["total"] < best_val_loss:
+            best_val_loss = val_losses["total"]
+            torch.save({
+                "epoch": epoch,
+                "model_state_dict": model.state_dict(),
+                "optimizer_state_dict": optimizer.state_dict(),
+                "val_loss": best_val_loss,
+                "config": config,
+            }, output_dir / "best_model.pt")
+            print(f"  ★ Best val loss: {best_val_loss:.4f}")
+
+        history.append({"epoch": epoch, "lr": lr, "train": train_losses, "val": val_losses})
+        if epoch % 20 == 0:
+            with open(output_dir / "history.json", "w") as f:
+                json.dump(history, f, indent=2)
+
+    torch.save({"epoch": args.epochs, "model_state_dict": model.state_dict(), "config": config},
+               output_dir / "final_model.pt")
+    with open(output_dir / "history.json", "w") as f:
+        json.dump(history, f, indent=2)
+
+    print(f"\nDone! Best val loss: {best_val_loss:.4f}")
+
+    # Push to Hub
+    if args.push_to_hub:
+        REPO_ID = args.push_to_hub
+        try:
+            create_repo(REPO_ID, repo_type="model", private=False, exist_ok=True)
+        except Exception:
+            pass
+        api = HfApi()
+        api.upload_file(path_or_fileobj=str(output_dir / "best_model.pt"),
+                        path_in_repo="best_model.pt", repo_id=REPO_ID, repo_type="model")
+        print(f"Model pushed to https://huggingface.co/{REPO_ID}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--image_dir", type=str, default="./images")
+    parser.add_argument("--label_dir", type=str, default="./pseudolabels")
+    parser.add_argument("--output_dir", type=str, default="./checkpoints_v3")
+    parser.add_argument("--base_ch", type=int, default=16)
+    parser.add_argument("--use_depthwise", action="store_true", default=True)
+    parser.add_argument("--no_depthwise", dest="use_depthwise", action="store_false")
+    parser.add_argument("--epochs", type=int, default=400)
+    parser.add_argument("--batch_size", type=int, default=4)
+    parser.add_argument("--lr", type=float, default=1e-3)
+    parser.add_argument("--weight_decay", type=float, default=1e-4)
+    parser.add_argument("--crop_size", type=int, default=256)
+    parser.add_argument("--val_split", type=float, default=0.15)
+    parser.add_argument("--flow_weight", type=float, default=1.0)
+    parser.add_argument("--prob_weight", type=float, default=1.0)
+    parser.add_argument("--dist_weight", type=float, default=1.0)
+    parser.add_argument("--no_gpu", action="store_true")
+    parser.add_argument("--push_to_hub", type=str, default=None)
+    args = parser.parse_args()
+    train(args)