Upload train_unet_advanced.py with huggingface_hub

train_unet_advanced.py

@@ -0,0 +1,242 @@
+"""
+train_unet_advanced.py
+Train MobileNetV2-UNet student with LR scheduler comparison (Phase 2, Step 5).
+
+Supports:
+  - Option 1: MobileNetV2 + U-Net skip connections (student_unet.py)
+  - Option 2: Cosine Annealing LR scheduler
+  - Option 3: Step Decay LR scheduler (halve every 20 epochs)
+  - Option 4: Knowledge distillation with best scheduler + U-Net
+  - MODE: 'standalone' or 'distill'
+
+Usage:
+  # Compare schedulers in standalone mode:
+  python train_unet_advanced.py   # edit SCHEDULER_TYPE below
+
+  # After finding best scheduler, run distillation:
+  # Set MODE='distill' and SCHEDULER_TYPE to the winner
+
+Requires pre-computed teacher outputs for distillation mode
+(run precompute_teacher.py first).
+"""
+import sys
+import os
+import time
+import torch
+import torch.nn as nn
+import numpy as np
+import warnings
+with warnings.catch_warnings():
+    warnings.filterwarnings("ignore", category=FutureWarning)
+
+from utils import localization_error, write_log
+from data_loader import load_data
+from student_unet import MobileNetUNetStudent
+
+# ============================================================
+# CONFIG
+# ============================================================
+MODE = 'standalone'           # 'standalone' or 'distill'
+SCHEDULER_TYPE = 'cosine'     # 'cosine' or 'step'
+
+EPOCHS = 120
+LR = 1e-4
+BATCH_SIZE = 32
+NUM_WORKERS = 4
+
+ALPHA = 0.7                   # ground truth loss weight (distillation)
+BETA = 0.3                    # teacher loss weight (distillation)
+
+TEACHER_TRAIN_PT = './data/teacher_train_outputs.pt'
+TEACHER_TEST_PT = './data/teacher_test_outputs.pt'
+
+# Data paths (Fig 10b -- Jacobs Hall)
+trainpath = [
+    './data/dataset_jacobs_July28.mat',
+    './data/dataset_non_fov_train_jacobs_July28_2.mat',
+    './data/dataset_fov_train_jacobs_July28_2.mat',
+]
+testpath = [
+    './data/dataset_fov_test_jacobs_July28_2.mat',
+    './data/dataset_non_fov_test_jacobs_July28_2.mat',
+]
+
+# ============================================================
+# SETUP
+# ============================================================
+torch.manual_seed(0)
+np.random.seed(0)
+
+device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+run_label = f'{MODE}_unet_{SCHEDULER_TYPE}'
+save_name = time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime())
+save_dir = os.path.join('./runs', f'{run_label}_{save_name}')
+os.makedirs(save_dir, exist_ok=True)
+
+print(f'Device: {device}')
+print(f'Mode: {MODE} | Scheduler: {SCHEDULER_TYPE} | Epochs: {EPOCHS}')
+print(f'Save dir: {save_dir}')
+
+# ============================================================
+# LOAD DATA
+# ============================================================
+print('Loading training data...')
+train_w_list, train_lbl_list = [], []
+for fp in trainpath:
+    _, w, lbl = load_data(fp)
+    if lbl.dim() == 3:
+        lbl = lbl.unsqueeze(1)
+    train_w_list.append(w)
+    train_lbl_list.append(lbl)
+train_w = torch.cat(train_w_list, dim=0)
+train_lbl = torch.cat(train_lbl_list, dim=0)
+print(f'Training samples: {train_w.shape[0]}')
+
+print('Loading testing data...')
+test_w_list, test_lbl_list = [], []
+for fp in testpath:
+    _, w, lbl = load_data(fp)
+    if lbl.dim() == 3:
+        lbl = lbl.unsqueeze(1)
+    test_w_list.append(w)
+    test_lbl_list.append(lbl)
+test_w = torch.cat(test_w_list, dim=0)
+test_lbl = torch.cat(test_lbl_list, dim=0)
+print(f'Testing samples: {test_w.shape[0]}')
+
+# Load teacher outputs if distilling
+if MODE == 'distill':
+    print('Loading pre-computed teacher outputs...')
+    train_teacher = torch.load(TEACHER_TRAIN_PT, weights_only=True)
+    test_teacher = torch.load(TEACHER_TEST_PT, weights_only=True)
+    assert train_teacher.shape[0] == train_w.shape[0], \
+        f'Teacher/data mismatch: {train_teacher.shape[0]} vs {train_w.shape[0]}'
+    print(f'Teacher outputs: train={train_teacher.shape}, test={test_teacher.shape}')
+    train_data = torch.utils.data.TensorDataset(train_w, train_lbl, train_teacher)
+    test_data = torch.utils.data.TensorDataset(test_w, test_lbl, test_teacher)
+else:
+    train_data = torch.utils.data.TensorDataset(train_w, train_lbl)
+    test_data = torch.utils.data.TensorDataset(test_w, test_lbl)
+
+train_loader = torch.utils.data.DataLoader(
+    train_data, batch_size=BATCH_SIZE, shuffle=True,
+    num_workers=NUM_WORKERS, pin_memory=True, persistent_workers=True,
+)
+test_loader = torch.utils.data.DataLoader(
+    test_data, batch_size=BATCH_SIZE, shuffle=False,
+    num_workers=NUM_WORKERS, pin_memory=True, persistent_workers=True,
+)
+
+# ============================================================
+# MODEL + OPTIMIZER + SCHEDULER
+# ============================================================
+student = MobileNetUNetStudent(input_nc=4).to(device)
+n_params = sum(p.numel() for p in student.parameters())
+print(f'U-Net student params: {n_params:,} ({n_params/1e6:.2f}M)')
+
+optimizer = torch.optim.Adam(student.parameters(), lr=LR)
+criterion = nn.MSELoss()
+
+if SCHEDULER_TYPE == 'cosine':
+    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=EPOCHS)
+elif SCHEDULER_TYPE == 'step':
+    # Halve LR every 20 epochs (professor's spec)
+    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.5)
+else:
+    raise ValueError(f'Unknown scheduler: {SCHEDULER_TYPE}')
+
+# ============================================================
+# TRAINING LOOP
+# ============================================================
+print(f'Starting {MODE} training with {SCHEDULER_TYPE} scheduler...')
+best_median = float('inf')
+
+for epoch in range(1, EPOCHS + 1):
+    epoch_start = time.time()
+
+    # --- TRAIN ---
+    student.train()
+    train_loss_sum = 0
+    train_errors = []
+
+    for data in train_loader:
+        inputs = data[0].to(device)
+        labels = data[1].to(device)
+
+        student_output = student(inputs)
+        L_GT = criterion(student_output, labels)
+
+        if MODE == 'distill':
+            teacher_out = data[2].to(device)
+            L_KD = criterion(student_output, teacher_out)
+            loss = ALPHA * L_GT + BETA * L_KD
+        else:
+            loss = L_GT
+
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+        train_loss_sum += loss.item()
+        train_errors.extend(localization_error(
+            student_output.detach().cpu().numpy(),
+            labels.cpu().numpy(),
+            scale=0.1,
+        ))
+
+    train_loss_avg = train_loss_sum / len(train_loader)
+    train_median = np.median(train_errors)
+
+    # --- TEST ---
+    student.eval()
+    test_loss_sum = 0
+    test_errors = []
+
+    with torch.no_grad():
+        for data in test_loader:
+            inputs = data[0].to(device)
+            labels = data[1].to(device)
+
+            student_output = student(inputs)
+            loss = criterion(student_output, labels)
+
+            test_loss_sum += loss.item()
+            test_errors.extend(localization_error(
+                student_output.cpu().numpy(),
+                labels.cpu().numpy(),
+                scale=0.1,
+            ))
+
+    test_loss_avg = test_loss_sum / len(test_loader)
+    test_median = np.median(test_errors)
+    test_p90 = np.percentile(test_errors, 90)
+    test_p99 = np.percentile(test_errors, 99)
+    current_lr = optimizer.param_groups[0]['lr']
+    elapsed = time.time() - epoch_start
+
+    # Step scheduler
+    scheduler.step()
+
+    # --- LOG ---
+    write_log([str(train_loss_avg)], run_label, log_dir=save_dir, log_type='train_loss')
+    write_log([str(train_median)], run_label, log_dir=save_dir, log_type='train_median_error')
+    write_log([str(test_loss_avg)], run_label, log_dir=save_dir, log_type='test_loss')
+    write_log([str(test_median)], run_label, log_dir=save_dir, log_type='test_median_error')
+    write_log([str(test_p90)], run_label, log_dir=save_dir, log_type='test_90_error')
+    write_log([str(test_p99)], run_label, log_dir=save_dir, log_type='test_99_error')
+
+    marker = ' <-- NEW BEST' if test_median < best_median else ''
+    print(f'Epoch {epoch}/{EPOCHS} | {elapsed:.0f}s | LR={current_lr:.6f} | '
+          f'train_med={train_median:.3f}m | '
+          f'test_med={test_median:.3f}m P90={test_p90:.3f}m P99={test_p99:.3f}m{marker}')
+
+    # --- SAVE ---
+    if test_median < best_median:
+        best_median = test_median
+        torch.save(student.state_dict(), os.path.join(save_dir, 'best_unet.pth'))
+
+    torch.save(student.state_dict(), os.path.join(save_dir, 'latest_unet.pth'))
+
+print(f'\nDone. Best median error: {best_median:.4f}m')
+print(f'Scheduler: {SCHEDULER_TYPE} | Mode: {MODE}')
+print(f'Results saved to: {save_dir}')