Upload train.py

train.py

@@ -1,621 +1,265 @@
 """
 PriviGaze Training Script - Privileged Distillation for Gaze Estimation
-
-Two-phase training:
-1. Teacher pre-training: Train teacher on privileged data (RGB eyes + blurred face)
-2. Student distillation: Train student with privileged distillation loss
-
-This script implements Phase 2 (distillation). Phase 1 (teacher pre-training)
-should be run first to produce a strong teacher model.
-
-Usage:
-    python train.py --mode distill --teacher-path ./teacher_best.pt --epochs 100
 """
-
-import os
-import sys
-import argparse
-import time
+import os, sys, argparse, time
 from pathlib import Path
 from collections import defaultdict
-
 import torch
-import torch.nn as nn
 from torch.optim import AdamW
-from torch.optim.lr_scheduler import CosineAnnealingLR, ReduceLROnPlateau
+from torch.optim.lr_scheduler import CosineAnnealingLR
 import numpy as np
 
-# Add parent directory to path
 sys.path.insert(0, str(Path(__file__).parent))
-
 from models.teacher import PriviGazeTeacher
 from models.student import PriviGazeStudent, count_parameters
-from models.distillation_loss import PriviGazeDistillationLoss
-from models.dataset import create_dataloaders, SyntheticGazeDataset
+from models.distillation_loss import PriviGazeDistillationLoss, L2CSLoss, AngularLoss
+from models.dataset import create_dataloaders
 
-# Trackio for experiment monitoring
 try:
-    import trackio
-    HAS_TRACKIO = True
+    import trackio; HAS_TRACKIO = True
 except ImportError:
-    HAS_TRACKIO = False
-    print("Warning: trackio not installed. Logging to stdout only.")
+    HAS_TRACKIO = False; print("Warning: trackio not installed.")
 
 
 class DistillationTrainer:
-    """Trains student model via privileged distillation from teacher."""
-    
-    def __init__(
-        self,
-        teacher: PriviGazeTeacher,
-        student: PriviGazeStudent,
-        distillation_loss: PriviGazeDistillationLoss,
-        train_loader,
-        val_loader,
-        device: torch.device,
-        lr: float = 1e-4,
-        weight_decay: float = 1e-4,
-        epochs: int = 100,
-        teacher_frozen: bool = True,
-        trackio_project: str = "privi-gaze",
-        trackio_run_name: str = "distill",
-    ):
+    def __init__(self, teacher, student, dist_loss, train_loader, val_loader,
+                 device, lr=1e-4, wd=1e-4, epochs=100, tproj="privi-gaze", trun="distill"):
         self.teacher = teacher.to(device)
         self.student = student.to(device)
-        self.distillation_loss = distillation_loss.to(device)
+        self.dist_loss = dist_loss.to(device)
         self.train_loader = train_loader
         self.val_loader = val_loader
         self.device = device
         self.epochs = epochs
-        self.trackio_project = trackio_project
-        self.trackio_run_name = trackio_run_name
-        
-        if teacher_frozen:
-            for param in self.teacher.parameters():
-                param.requires_grad = False
-            self.teacher.eval()
-        
-        # Optimizer: only student parameters
-        self.optimizer = AdamW(
-            self.student.parameters(),
-            lr=lr,
-            weight_decay=weight_decay,
-        )
-        
-        # Scheduler
-        self.scheduler = CosineAnnealingLR(
-            self.optimizer,
-            T_max=epochs,
-            eta_min=lr * 0.01,
-        )
-        
-        # Track best model
-        self.best_val_loss = float('inf')
+        for p in self.teacher.parameters(): p.requires_grad = False
+        self.teacher.eval()
+        self.opt = AdamW(self.student.parameters(), lr=lr, weight_decay=wd)
+        self.sched = CosineAnnealingLR(self.opt, T_max=epochs, eta_min=lr*0.01)
+        self.best_val = float('inf')
         self.best_epoch = 0
-        
-        # Metrics tracking
-        self.metrics_history = defaultdict(list)
-        
-        # Initialize trackio
+        self.metrics = defaultdict(list)
         if HAS_TRACKIO:
-            trackio.init(
-                project=trackio_project,
-                run_name=trackio_run_name,
-                config={
-                    'student_params': count_parameters(self.student),
-                    'teacher_params': count_parameters(self.teacher),
-                    'lr': lr,
-                    'weight_decay': weight_decay,
-                    'epochs': epochs,
-                    'batch_size': train_loader.batch_size,
-                }
-            )
-    
-    def train_epoch(self, epoch: int) -> dict:
-        """Train for one epoch."""
+            trackio.init(project=tproj, run_name=trun,
+                config={'student_params': count_parameters(student),
+                        'teacher_params': count_parameters(teacher), 'lr': lr, 'epochs': epochs})
+
+    def train_epoch(self, epoch):
         self.student.train()
-        epoch_losses = defaultdict(float)
-        num_batches = 0
-        
-        for batch_idx, batch in enumerate(self.train_loader):
-            # Move to device
-            left_eye = batch['left_eye'].to(self.device)
-            right_eye = batch['right_eye'].to(self.device)
-            face_blurred = batch['face_blurred_gray'].to(self.device)
-            face_gray = batch['face_gray'].to(self.device)
-            pitch_target = batch['pitch'].to(self.device)
-            yaw_target = batch['yaw'].to(self.device)
-            
-            # Teacher forward (no grad)
+        losses = defaultdict(float)
+        n = 0
+        for bi, batch in enumerate(self.train_loader):
+            le = batch['left_eye'].to(self.device)
+            re = batch['right_eye'].to(self.device)
+            fb = batch['face_blurred_gray'].to(self.device)
+            fg = batch['face_gray'].to(self.device)
+            pt = batch['pitch'].to(self.device)
+            yt = batch['yaw'].to(self.device)
             with torch.no_grad():
-                t_pitch, t_yaw, t_features = self.teacher(
-                    left_eye, right_eye, face_blurred
-                )
-                # Get teacher logits by running forward through heads
-                # (We need these for logit distillation)
-                # We extract them from the teacher's internal state
-                t_pitch_logits = self.teacher.pitch_head(t_features)
-                t_yaw_logits = self.teacher.yaw_head(t_features)
-            
-            # Student forward
-            s_pitch, s_yaw, s_features = self.student(face_gray)
-            s_pitch_logits = self.student.pitch_head(s_features)
-            s_yaw_logits = self.student.yaw_head(s_features)
-            
-            # Compute distillation loss
-            loss, loss_dict = self.distillation_loss(
-                s_pitch, s_yaw,
-                s_pitch_logits, s_yaw_logits,
-                s_features,
-                t_pitch, t_yaw,
-                t_pitch_logits, t_yaw_logits,
-                t_features,
-                pitch_target, yaw_target,
-            )
-            
-            # Backward
-            self.optimizer.zero_grad()
+                tp, ty, tplog, tylog, tf = self.teacher(le, re, fb)
+            sp, sy, sf = self.student(fg)
+            splog = self.student.pitch_head(sf)
+            sylog = self.student.yaw_head(sf)
+            loss, ld = self.dist_loss(sp, sy, splog, sylog, sf,
+                                       tp, ty, tplog, tylog, tf, pt, yt)
+            self.opt.zero_grad()
             loss.backward()
-            
-            # Gradient clipping
-            torch.nn.utils.clip_grad_norm_(self.student.parameters(), max_norm=1.0)
-            
-            self.optimizer.step()
-            
-            # Accumulate losses
-            for k, v in loss_dict.items():
-                epoch_losses[k] += v
-            num_batches += 1
-            
-            # Log every 100 batches
-            if batch_idx % 100 == 0:
-                self._log_step(epoch, batch_idx, loss_dict)
-        
-        # Average losses
-        for k in epoch_losses:
-            epoch_losses[k] /= num_batches
-        
-        return dict(epoch_losses)
-    
+            torch.nn.utils.clip_grad_norm_(self.student.parameters(), 1.0)
+            self.opt.step()
+            for k, v in ld.items(): losses[k] += v
+            n += 1
+            if bi % 100 == 0:
+                print(f"Epoch {epoch} | Batch {bi} | " + " | ".join(f"{k}={v:.4f}" for k, v in ld.items()))
+                if HAS_TRACKIO:
+                    for k2, v2 in ld.items(): trackio.log({f"train/{k2}": v2})
+        return {k: v/n for k, v in losses.items()}
+
     @torch.no_grad()
-    def validate(self, epoch: int) -> dict:
-        """Validate the student model."""
+    def validate(self, epoch):
         self.student.eval()
         self.teacher.eval()
-        
-        val_losses = defaultdict(float)
-        angular_errors = []
-        pitch_errors = []
-        yaw_errors = []
-        num_batches = 0
-        
+        losses = defaultdict(float)
+        ae, pe, ye = [], [], []
+        n = 0
         for batch in self.val_loader:
-            left_eye = batch['left_eye'].to(self.device)
-            right_eye = batch['right_eye'].to(self.device)
-            face_blurred = batch['face_blurred_gray'].to(self.device)
-            face_gray = batch['face_gray'].to(self.device)
-            pitch_target = batch['pitch'].to(self.device)
-            yaw_target = batch['yaw'].to(self.device)
-            
-            # Teacher forward
-            t_pitch, t_yaw, t_features = self.teacher(
-                left_eye, right_eye, face_blurred
-            )
-            t_pitch_logits = self.teacher.pitch_head(t_features)
-            t_yaw_logits = self.teacher.yaw_head(t_features)
-            
-            # Student forward
-            s_pitch, s_yaw, s_features = self.student(face_gray)
-            s_pitch_logits = self.student.pitch_head(s_features)
-            s_yaw_logits = self.student.yaw_head(s_features)
-            
-            # Compute loss
-            loss, loss_dict = self.distillation_loss(
-                s_pitch, s_yaw,
-                s_pitch_logits, s_yaw_logits,
-                s_features,
-                t_pitch, t_yaw,
-                t_pitch_logits, t_yaw_logits,
-                t_features,
-                pitch_target, yaw_target,
-            )
-            
-            for k, v in loss_dict.items():
-                val_losses[k] += v
-            num_batches += 1
-            
-            # Compute angular error
-            angular_err = torch.sqrt(
-                (s_pitch - pitch_target) ** 2 + (s_yaw - yaw_target) ** 2
-            )
-            angular_errors.extend(angular_err.cpu().tolist())
-            pitch_errors.extend((s_pitch - pitch_target).abs().cpu().tolist())
-            yaw_errors.extend((s_yaw - yaw_target).abs().cpu().tolist())
-        
-        for k in val_losses:
-            val_losses[k] /= num_batches
-        
-        val_losses['angular_error_mean'] = np.mean(angular_errors)
-        val_losses['angular_error_std'] = np.std(angular_errors)
-        val_losses['pitch_error_mean'] = np.mean(pitch_errors)
-        val_losses['yaw_error_mean'] = np.mean(yaw_errors)
-        
-        return dict(val_losses)
-    
-    def _log_step(self, epoch, batch_idx, loss_dict):
-        """Log training step metrics."""
-        msg = f"Epoch {epoch} | Batch {batch_idx} | "
-        msg += " | ".join(f"{k}={v:.4f}" for k, v in loss_dict.items())
-        print(msg)
-        
-        if HAS_TRACKIO:
-            for k, v in loss_dict.items():
-                trackio.log({f"train/{k}": v})
-    
-    def _log_epoch(self, epoch, train_losses, val_losses):
-        """Log epoch metrics."""
-        print(f"\n{'='*60}")
-        print(f"Epoch {epoch} Summary:")
-        print(f"  Train: ", " | ".join(f"{k}={v:.4f}" for k, v in train_losses.items()))
-        print(f"  Val:   ", " | ".join(f"{k}={v:.4f}" for k, v in val_losses.items()))
-        print(f"{'='*60}\n")
-        
-        if HAS_TRACKIO:
-            for k, v in train_losses.items():
-                trackio.log({f"epoch/train_{k}": v}, step=epoch)
-            for k, v in val_losses.items():
-                trackio.log({f"epoch/val_{k}": v}, step=epoch)
-            
-            # Alert on overfitting
-            if epoch > 10 and val_losses.get('loss_total', 0) > self.best_val_loss * 1.3:
-                trackio.alert(
-                    "Possible Overfitting",
-                    f"Val loss {val_losses['loss_total']:.4f} >> best {self.best_val_loss:.4f} at epoch {epoch}",
-                    level="WARN",
-                )
-    
-    def train(self, save_dir: str = "./checkpoints"):
-        """Full training loop."""
+            le = batch['left_eye'].to(self.device)
+            re = batch['right_eye'].to(self.device)
+            fb = batch['face_blurred_gray'].to(self.device)
+            fg = batch['face_gray'].to(self.device)
+            pt = batch['pitch'].to(self.device)
+            yt = batch['yaw'].to(self.device)
+            tp, ty, tplog, tylog, tf = self.teacher(le, re, fb)
+            sp, sy, sf = self.student(fg)
+            splog = self.student.pitch_head(sf)
+            sylog = self.student.yaw_head(sf)
+            loss, ld = self.dist_loss(sp, sy, splog, sylog, sf,
+                                       tp, ty, tplog, tylog, tf, pt, yt)
+            for k, v in ld.items(): losses[k] += v
+            n += 1
+            aerr = torch.sqrt((sp-pt)**2 + (sy-yt)**2)
+            ae.extend(aerr.cpu().tolist())
+            pe.extend((sp-pt).abs().cpu().tolist())
+            ye.extend((sy-yt).abs().cpu().tolist())
+        for k in losses: losses[k] /= n
+        losses['angular_mean'] = np.mean(ae)
+        losses['angular_std'] = np.std(ae)
+        losses['pitch_mean'] = np.mean(pe)
+        losses['yaw_mean'] = np.mean(ye)
+        return dict(losses)
+
+    def train(self, save_dir="./checkpoints"):
         os.makedirs(save_dir, exist_ok=True)
-        
-        print(f"Starting distillation training for {self.epochs} epochs")
-        print(f"Student parameters: {count_parameters(self.student):,}")
-        print(f"Device: {self.device}")
-        
-        start_time = time.time()
-        
+        print(f"Distillation: {self.epochs} epochs | Student: {count_parameters(self.student):,} params")
+        t0 = time.time()
         for epoch in range(self.epochs):
-            epoch_start = time.time()
-            
-            # Train
-            train_losses = self.train_epoch(epoch)
-            
-            # Validate
-            val_losses = self.validate(epoch)
-            
-            # Step scheduler
-            self.scheduler.step()
-            current_lr = self.optimizer.param_groups[0]['lr']
-            
-            # Log
-            self._log_epoch(epoch, train_losses, val_losses)
-            
-            # Track metrics
-            for k, v in train_losses.items():
-                self.metrics_history[f'train_{k}'].append(v)
-            for k, v in val_losses.items():
-                self.metrics_history[f'val_{k}'].append(v)
-            
-            # Save best model
-            val_total = val_losses.get('loss_total', val_losses.get('angular_error_mean', float('inf')))
-            if val_total < self.best_val_loss:
-                self.best_val_loss = val_total
+            te = time.time()
+            tl = self.train_epoch(epoch)
+            vl = self.validate(epoch)
+            self.sched.step()
+            lr = self.opt.param_groups[0]['lr']
+            print(f"\n{'='*60}")
+            print(f"Epoch {epoch}: train={tl.get('loss_total',0):.4f} val={vl.get('loss_total',0):.4f} angular={vl.get('angular_mean',0):.2f}deg")
+            print(f"{'='*60}\n")
+            for k, v in tl.items(): self.metrics[f'train_{k}'].append(v)
+            for k, v in vl.items(): self.metrics[f'val_{k}'].append(v)
+            vt = vl.get('loss_total', vl.get('angular_mean', float('inf')))
+            if vt < self.best_val:
+                self.best_val = vt
                 self.best_epoch = epoch
-                
-                torch.save({
-                    'epoch': epoch,
-                    'student_state_dict': self.student.state_dict(),
-                    'optimizer_state_dict': self.optimizer.state_dict(),
-                    'best_val_loss': self.best_val_loss,
-                    'metrics_history': dict(self.metrics_history),
-                }, os.path.join(save_dir, 'student_best.pt'))
-                
-                if HAS_TRACKIO:
-                    trackio.alert(
-                        "New Best Model",
-                        f"Val loss: {val_total:.4f} at epoch {epoch} (angular: {val_losses.get('angular_error_mean', 0):.2f}°)",
-                        level="INFO",
-                    )
-            
-            # Save checkpoint every 10 epochs
+                torch.save({'epoch': epoch, 'student_state_dict': self.student.state_dict(),
+                            'opt_state_dict': self.opt.state_dict(), 'best_val': self.best_val,
+                            'metrics': dict(self.metrics)}, os.path.join(save_dir, 'student_best.pt'))
+                if HAS_TRACKIO: trackio.alert("New Best", f"Val {vt:.4f} @ epoch {epoch}", level="INFO")
             if epoch % 10 == 0:
-                torch.save({
-                    'epoch': epoch,
-                    'student_state_dict': self.student.state_dict(),
-                    'optimizer_state_dict': self.optimizer.state_dict(),
-                }, os.path.join(save_dir, f'student_epoch_{epoch}.pt'))
-            
-            epoch_time = time.time() - epoch_start
-            print(f"Epoch {epoch} took {epoch_time:.1f}s, LR: {current_lr:.2e}")
-        
-        total_time = time.time() - start_time
-        print(f"\nTraining complete! Total time: {total_time/3600:.1f}h")
-        print(f"Best validation loss: {self.best_val_loss:.4f} at epoch {self.best_epoch}")
-        
-        if HAS_TRACKIO:
-            trackio.alert(
-                "Training Complete",
-                f"Best val loss: {self.best_val_loss:.4f} at epoch {self.best_epoch}. "
-                f"Student params: {count_parameters(self.student):,}",
-                level="INFO",
-            )
-        
-        return self.best_val_loss
+                torch.save({'epoch': epoch, 'student_state_dict': self.student.state_dict(),
+                            'opt_state_dict': self.opt.state_dict()},
+                           os.path.join(save_dir, f'student_epoch_{epoch}.pt'))
+            print(f"Epoch {epoch} took {time.time()-te:.1f}s, LR={lr:.2e}")
+        print(f"\nDone! Best val: {self.best_val:.4f} @ epoch {self.best_epoch}")
+        return self.best_val
 
 
-def pretrain_teacher(
-    teacher: PriviGazeTeacher,
-    train_loader,
-    val_loader,
-    device: torch.device,
-    lr: float = 1e-4,
-    epochs: int = 50,
-    save_dir: str = "./checkpoints",
-) -> str:
-    """Pre-train the teacher model on privileged data."""
-    from models.distillation_loss import L2CSLoss, AngularLoss
-    
+def pretrain_teacher(teacher, train_loader, val_loader, device, lr=1e-4, epochs=50, save_dir="./checkpoints"):
     teacher = teacher.to(device)
-    teacher.train()
-    
-    optimizer = AdamW(teacher.parameters(), lr=lr, weight_decay=1e-4)
-    scheduler = CosineAnnealingLR(optimizer, T_max=epochs, eta_min=lr * 0.01)
-    
-    pitch_loss_fn = L2CSLoss(gaze_bins=90)
-    yaw_loss_fn = L2CSLoss(gaze_bins=90)
-    angular_loss_fn = AngularLoss()
-    
-    best_val_loss = float('inf')
+    opt = AdamW(teacher.parameters(), lr=lr, weight_decay=1e-4)
+    sched = CosineAnnealingLR(opt, T_max=epochs, eta_min=lr*0.01)
+    ploss = L2CSLoss(gaze_bins=90)
+    yloss = L2CSLoss(gaze_bins=90)
+    aloss = AngularLoss()
+    best = float('inf')
     os.makedirs(save_dir, exist_ok=True)
-    
     for epoch in range(epochs):
-        # Training
         teacher.train()
-        train_loss_total = 0.0
+        tloss = 0.0
         for batch in train_loader:
-            left_eye = batch['left_eye'].to(device)
-            right_eye = batch['right_eye'].to(device)
-            face_blurred = batch['face_blurred_gray'].to(device)
-            pitch_target = batch['pitch'].to(device)
-            yaw_target = batch['yaw'].to(device)
-            
-            pitch_pred, yaw_pred, features = teacher(left_eye, right_eye, face_blurred)
-            pitch_logits = teacher.pitch_head(features)
-            yaw_logits = teacher.yaw_head(features)
-            
-            loss = (pitch_loss_fn(pitch_logits, pitch_pred, pitch_target) +
-                    yaw_loss_fn(yaw_logits, yaw_pred, yaw_target) +
-                    angular_loss_fn(pitch_pred, yaw_pred, pitch_target, yaw_target))
-            
-            optimizer.zero_grad()
+            le = batch['left_eye'].to(device)
+            re = batch['right_eye'].to(device)
+            fb = batch['face_blurred_gray'].to(device)
+            pt = batch['pitch'].to(device)
+            yt = batch['yaw'].to(device)
+            pp, yp, pl, yl, _ = teacher(le, re, fb)
+            loss = ploss(pl, pp, pt) + yloss(yl, yp, yt) + aloss(pp, yp, pt, yt)
+            opt.zero_grad()
             loss.backward()
-            torch.nn.utils.clip_grad_norm_(teacher.parameters(), max_norm=1.0)
-            optimizer.step()
-            
-            train_loss_total += loss.item()
-        
-        train_loss_total /= len(train_loader)
-        
-        # Validation
+            torch.nn.utils.clip_grad_norm_(teacher.parameters(), 1.0)
+            opt.step()
+            tloss += loss.item()
+        tloss /= len(train_loader)
         teacher.eval()
-        val_loss_total = 0.0
-        val_angular = 0.0
+        vloss = 0.0
+        va = 0.0
         with torch.no_grad():
             for batch in val_loader:
-                left_eye = batch['left_eye'].to(device)
-                right_eye = batch['right_eye'].to(device)
-                face_blurred = batch['face_blurred_gray'].to(device)
-                pitch_target = batch['pitch'].to(device)
-                yaw_target = batch['yaw'].to(device)
-                
-                pitch_pred, yaw_pred, features = teacher(left_eye, right_eye, face_blurred)
-                pitch_logits = teacher.pitch_head(features)
-                yaw_logits = teacher.yaw_head(features)
-                
-                loss = (pitch_loss_fn(pitch_logits, pitch_pred, pitch_target) +
-                        yaw_loss_fn(yaw_logits, yaw_pred, yaw_target))
-                val_loss_total += loss.item()
-                
-                angular_err = torch.sqrt((pitch_pred - pitch_target)**2 + (yaw_pred - yaw_target)**2)
-                val_angular += angular_err.mean().item()
-        
-        val_loss_total /= len(val_loader)
-        val_angular /= len(val_loader)
-        
-        scheduler.step()
-        
-        print(f"Teacher Epoch {epoch}: train_loss={train_loss_total:.4f}, "
-              f"val_loss={val_loss_total:.4f}, val_angular={val_angular:.2f}°")
-        
-        if val_loss_total < best_val_loss:
-            best_val_loss = val_loss_total
+                le = batch['left_eye'].to(device)
+                re = batch['right_eye'].to(device)
+                fb = batch['face_blurred_gray'].to(device)
+                pt = batch['pitch'].to(device)
+                yt = batch['yaw'].to(device)
+                pp, yp, pl, yl, _ = teacher(le, re, fb)
+                vloss += (ploss(pl, pp, pt) + yloss(yl, yp, yt)).item()
+                va += torch.sqrt((pp-pt)**2 + (yp-yt)**2).mean().item()
+        vloss /= len(val_loader)
+        va /= len(val_loader)
+        sched.step()
+        print(f"Teacher Epoch {epoch}: train={tloss:.4f} val={vloss:.4f} angular={va:.2f}deg")
+        if vloss < best:
+            best = vloss
             torch.save(teacher.state_dict(), os.path.join(save_dir, 'teacher_best.pt'))
-    
     return os.path.join(save_dir, 'teacher_best.pt')
 
 
 def main():
-    parser = argparse.ArgumentParser(description="PriviGaze Distillation Training")
-    parser.add_argument('--mode', type=str, default='distill',
-                       choices=['pretrain_teacher', 'distill', 'both'],
-                       help='Training mode')
-    parser.add_argument('--teacher-path', type=str, default=None,
-                       help='Path to pre-trained teacher checkpoint')
-    parser.add_argument('--batch-size', type=int, default=32,
-                       help='Batch size')
-    parser.add_argument('--epochs', type=int, default=100,
-                       help='Number of distillation epochs')
-    parser.add_argument('--teacher-epochs', type=int, default=50,
-                       help='Number of teacher pre-training epochs')
-    parser.add_argument('--lr', type=float, default=1e-4,
-                       help='Learning rate')
-    parser.add_argument('--weight-decay', type=float, default=1e-4,
-                       help='Weight decay')
-    parser.add_argument('--num-train', type=int, default=40000,
-                       help='Number of synthetic training samples')
-    parser.add_argument('--num-val', type=int, default=5000,
-                       help='Number of synthetic val samples')
-    parser.add_argument('--save-dir', type=str, default='./checkpoints',
-                       help='Directory to save checkpoints')
-    parser.add_argument('--device', type=str, default='cuda',
-                       help='Device to train on')
-    parser.add_argument('--trackio-project', type=str, default='privi-gaze',
-                       help='Trackio project name')
-    parser.add_argument('--trackio-run', type=str, default='distill-run',
-                       help='Trackio run name')
-    parser.add_argument('--push-to-hub', action='store_true',
-                       help='Push trained model to HF Hub')
-    parser.add_argument('--hub-model-id', type=str, default=None,
-                       help='HF Hub model ID for pushing')
-    parser.add_argument('--alpha-contrastive', type=float, default=0.5,
-                       help='Weight for contrastive distillation loss')
-    parser.add_argument('--alpha-mmd', type=float, default=0.1,
-                       help='Weight for MMD distribution matching loss')
-    parser.add_argument('--alpha-logit', type=float, default=0.5,
-                       help='Weight for logit distillation loss')
-    
-    args = parser.parse_args()
-    
-    # Device setup
+    p = argparse.ArgumentParser(description="PriviGaze Training")
+    p.add_argument('--mode', type=str, default='distill', choices=['pretrain_teacher','distill','both'])
+    p.add_argument('--teacher-path', type=str, default=None)
+    p.add_argument('--batch-size', type=int, default=32)
+    p.add_argument('--epochs', type=int, default=100)
+    p.add_argument('--teacher-epochs', type=int, default=50)
+    p.add_argument('--lr', type=float, default=1e-4)
+    p.add_argument('--weight-decay', type=float, default=1e-4)
+    p.add_argument('--num-train', type=int, default=40000)
+    p.add_argument('--num-val', type=int, default=5000)
+    p.add_argument('--save-dir', type=str, default='./checkpoints')
+    p.add_argument('--device', type=str, default='cuda')
+    p.add_argument('--trackio-project', type=str, default='privi-gaze')
+    p.add_argument('--trackio-run', type=str, default='distill-run')
+    p.add_argument('--push-to-hub', action='store_true')
+    p.add_argument('--hub-model-id', type=str, default=None)
+    p.add_argument('--alpha-contrastive', type=float, default=0.5)
+    p.add_argument('--alpha-mmd', type=float, default=0.1)
+    p.add_argument('--alpha-logit', type=float, default=0.5)
+    args = p.parse_args()
+
     device = torch.device(args.device if torch.cuda.is_available() else 'cpu')
-    print(f"Using device: {device}")
-    
-    # Create dataloaders
+    print(f"Device: {device}")
     train_loader, val_loader, test_loader = create_dataloaders(
-        num_train=args.num_train,
-        num_val=args.num_val,
-        batch_size=args.batch_size,
-    )
-    
-    # Initialize models
-    teacher = PriviGazeTeacher(
-        eye_backbone="facebook/convnextv2-atto-1k-224",
-        face_backbone="facebook/convnextv2-nano-22k-384",
-    )
-    
+        num_train=args.num_train, num_val=args.num_val, batch_size=args.batch_size)
+
+    teacher = PriviGazeTeacher()
     student = PriviGazeStudent()
-    
-    print(f"Teacher parameters: {count_parameters(teacher):,}")
-    print(f"Student parameters: {count_parameters(student):,}")
-    
-    # Pre-train teacher if needed
+    print(f"Teacher: {count_parameters(teacher):,} params")
+    print(f"Student: {count_parameters(student):,} params")
+
     if args.mode in ['pretrain_teacher', 'both']:
-        print("\n=== Phase 1: Pre-training Teacher ===")
-        teacher_path = pretrain_teacher(
-            teacher, train_loader, val_loader, device,
-            lr=args.lr, epochs=args.teacher_epochs,
-            save_dir=args.save_dir,
-        )
-        print(f"Teacher saved to: {teacher_path}")
-        args.teacher_path = teacher_path
-    
-    # Load teacher checkpoint
+        print("\n=== Phase 1: Teacher Pre-training ===")
+        tp = pretrain_teacher(teacher, train_loader, val_loader, device,
+                              lr=args.lr, epochs=args.teacher_epochs, save_dir=args.save_dir)
+        args.teacher_path = tp
+
     if args.teacher_path:
-        print(f"\nLoading teacher from: {args.teacher_path}")
+        print(f"\nLoading teacher: {args.teacher_path}")
         teacher.load_state_dict(torch.load(args.teacher_path, map_location=device))
-    
-    # Distill
+
     if args.mode in ['distill', 'both']:
-        print("\n=== Phase 2: Privileged Distillation ===")
-        
-        # Create distillation loss
-        dist_loss = PriviGazeDistillationLoss(
-            gaze_bins=90,
-            teacher_feature_dim=256,
-            student_feature_dim=128,
-            alpha_contrastive=args.alpha_contrastive,
-            alpha_mmd=args.alpha_mmd,
-            alpha_logit=args.alpha_logit,
-        )
-        
-        # Create trainer
-        trainer = DistillationTrainer(
-            teacher=teacher,
-            student=student,
-            distillation_loss=dist_loss,
-            train_loader=train_loader,
-            val_loader=val_loader,
-            device=device,
-            lr=args.lr,
-            weight_decay=args.weight_decay,
-            epochs=args.epochs,
-            trackio_project=args.trackio_project,
-            trackio_run_name=args.trackio_run,
-        )
-        
-        # Train
-        best_loss = trainer.train(save_dir=args.save_dir)
-        
-        # Test evaluation
-        print("\n=== Final Test Evaluation ===")
-        student.eval()
-        student.to(device)
-        
-        test_angular_errors = []
+        print("\n=== Phase 2: Distillation ===")
+        dloss = PriviGazeDistillationLoss(
+            gaze_bins=90, teacher_feature_dim=256, student_feature_dim=128,
+            alpha_contrastive=args.alpha_contrastive, alpha_mmd=args.alpha_mmd,
+            alpha_logit=args.alpha_logit)
+        trainer = DistillationTrainer(teacher, student, dloss, train_loader, val_loader,
+            device, lr=args.lr, wd=args.weight_decay, epochs=args.epochs,
+            tproj=args.trackio_project, trun=args.trackio_run)
+        trainer.train(save_dir=args.save_dir)
+
+        print("\n=== Test ===")
+        student.eval().to(device)
+        terr = []
         with torch.no_grad():
             for batch in test_loader:
-                face_gray = batch['face_gray'].to(device)
-                pitch_target = batch['pitch'].to(device)
-                yaw_target = batch['yaw'].to(device)
-                
-                pitch_pred, yaw_pred, _ = student(face_gray)
-                
-                angular_err = torch.sqrt(
-                    (pitch_pred - pitch_target) ** 2 + (yaw_pred - yaw_target) ** 2
-                )
-                test_angular_errors.extend(angular_err.cpu().tolist())
-        
-        mean_error = np.mean(test_angular_errors)
-        std_error = np.std(test_angular_errors)
-        print(f"Test Angular Error: {mean_error:.2f}° ± {std_error:.2f}°")
-        
-        if HAS_TRACKIO:
-            trackio.log({
-                'test/angular_error_mean': mean_error,
-                'test/angular_error_std': std_error,
-            })
-            trackio.alert(
-                "Test Results",
-                f"Angular error: {mean_error:.2f}° ± {std_error:.2f}°. "
-                f"Student params: {count_parameters(student):,}",
-                level="INFO",
-            )
-        
-        # Push to hub
+                fg = batch['face_gray'].to(device)
+                pt = batch['pitch'].to(device)
+                yt = batch['yaw'].to(device)
+                sp, sy, _ = student(fg)
+                terr.extend(torch.sqrt((sp-pt)**2 + (sy-yt)**2).cpu().tolist())
+        me = np.mean(terr); se = np.std(terr)
+        print(f"Test Angular Error: {me:.2f}deg +- {se:.2f}deg")
+
         if args.push_to_hub and args.hub_model_id:
             from huggingface_hub import HfApi
-            api = HfApi()
-            
-            # Save final model
-            model_path = os.path.join(args.save_dir, 'student_final.pt')
-            torch.save({
-                'student_state_dict': student.state_dict(),
-                'config': {
-                    'params': count_parameters(student),
-                    'test_angular_error': mean_error,
-                }
-            }, model_path)
-            
-            # Upload
-            api.upload_file(
-                path_or_fileobj=model_path,
-                path_in_repo="student_model.pt",
-                repo_id=args.hub_model_id,
-            )
-            print(f"Model pushed to: https://huggingface.co/{args.hub_model_id}")
-    
-    return best_loss if args.mode in ['distill', 'both'] else None
-
+            mp = os.path.join(args.save_dir, 'student_final.pt')
+            torch.save({'student_state_dict': student.state_dict(),
+                        'config': {'params': count_parameters(student), 'test_err': me}}, mp)
+            HfApi().upload_file(path_or_fileobj=mp, path_in_repo="student_model.pt", repo_id=args.hub_model_id)
+            print(f"Pushed to: https://huggingface.co/{args.hub_model_id}")
 
 if __name__ == "__main__":
     main()