Upload scripts/train_distillation.py with huggingface_hub

scripts/train_distillation.py

@@ -0,0 +1,239 @@
+#!/usr/bin/python
+"""
+train_distillation.py
+Knowledge distillation: DLoc teacher -> lightweight student.
+
+Loss: L = ALPHA * MSE(student, ground_truth) + BETA * MSE(student, teacher_output)
+Default: ALPHA=0.7, BETA=0.3 (from professor's spec)
+
+Requires pre-computed teacher outputs (run precompute_teacher.py first).
+Set MODE='standalone' to train without distillation (ground truth only).
+
+Usage:
+  python precompute_teacher.py          # run once
+  python train_distillation.py          # distillation (default: mobilenet)
+
+  To change student or mode, edit CONFIG below.
+"""
+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_mobilenet import MobileNetStudent
+from student_tinycnn import TinyCNNStudent
+
+# ============================================================
+# CONFIG — edit these before running
+# ============================================================
+STUDENT_TYPE = 'tinycnn'
+MODE = 'distill'
+
+ALPHA = 0.7                         # ground truth loss weight
+BETA = 0.3                          # distillation loss weight (ignored in standalone)
+EPOCHS = 50
+BATCH_SIZE = 32
+LR = 1e-4
+NUM_WORKERS = 4
+
+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')
+save_name = time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime())
+run_label = f'{MODE}_{STUDENT_TYPE}'
+save_dir = os.path.join('./runs', f'{run_label}_{save_name}')
+os.makedirs(save_dir, exist_ok=True)
+
+print(f'Mode: {MODE}')
+print(f'Student: {STUDENT_TYPE}')
+if MODE == 'distill':
+    print(f'Loss: {ALPHA}*L_GT + {BETA}*L_KD')
+else:
+    print(f'Loss: L_GT (standalone, no distillation)')
+print(f'Save dir: {save_dir}')
+
+# ============================================================
+# LOAD DATA INTO RAM
+# ============================================================
+print('Loading training data...')
+train_wo, train_w, train_lbl = [], [], []
+for fp in trainpath:
+    wo, w, lbl = load_data(fp)
+    if lbl.dim() == 3:
+        lbl = lbl.unsqueeze(1)
+    train_wo.append(wo)
+    train_w.append(w)
+    train_lbl.append(lbl)
+train_w = torch.cat(train_w, dim=0)
+train_lbl = torch.cat(train_lbl, dim=0)
+print(f'Training samples: {train_w.shape[0]}')
+
+print('Loading testing data...')
+test_wo, test_w, test_lbl = [], [], []
+for fp in testpath:
+    wo, w, lbl = load_data(fp)
+    if lbl.dim() == 3:
+        lbl = lbl.unsqueeze(1)
+    test_wo.append(wo)
+    test_w.append(w)
+    test_lbl.append(lbl)
+test_w = torch.cat(test_w, dim=0)
+test_lbl = torch.cat(test_lbl, dim=0)
+print(f'Testing samples: {test_w.shape[0]}')
+
+# Load pre-computed 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 loaded: 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,
+)
+
+# ============================================================
+# CREATE STUDENT
+# ============================================================
+if STUDENT_TYPE == 'mobilenet':
+    student = MobileNetStudent(input_nc=4).to(device)
+elif STUDENT_TYPE == 'tinycnn':
+    student = TinyCNNStudent(input_nc=4).to(device)
+else:
+    raise ValueError(f'Unknown student type: {STUDENT_TYPE}')
+
+n_params = sum(p.numel() for p in student.parameters())
+print(f'Student params: {n_params:,} ({n_params/1e6:.2f}M)')
+
+optimizer = torch.optim.Adam(student.parameters(), lr=LR)
+criterion = nn.MSELoss()
+
+# ============================================================
+# TRAINING LOOP
+# ============================================================
+print(f'Starting {MODE} training...')
+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)
+    elapsed = time.time() - epoch_start
+
+    # --- 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')
+
+    print(f'Epoch {epoch}/{EPOCHS} | {elapsed:.0f}s | '
+          f'train_loss={train_loss_avg:.6f} train_med={train_median:.3f}m | '
+          f'test_med={test_median:.3f}m P90={test_p90:.3f}m P99={test_p99:.3f}m')
+
+    # --- SAVE ---
+    if test_median < best_median:
+        best_median = test_median
+        torch.save(student.state_dict(), os.path.join(save_dir, f'best_{STUDENT_TYPE}.pth'))
+        print(f'  -> New best: {best_median:.4f}m')
+
+    torch.save(student.state_dict(), os.path.join(save_dir, f'latest_{STUDENT_TYPE}.pth'))
+
+print(f'\nDone. Best median error: {best_median:.4f}m')
+print(f'Results saved to: {save_dir}')