Upload train.py

train.py

@@ -0,0 +1,621 @@
+"""
+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
+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
+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
+
+# Trackio for experiment monitoring
+try:
+    import trackio
+    HAS_TRACKIO = True
+except ImportError:
+    HAS_TRACKIO = False
+    print("Warning: trackio not installed. Logging to stdout only.")
+
+
+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",
+    ):
+        self.teacher = teacher.to(device)
+        self.student = student.to(device)
+        self.distillation_loss = distillation_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')
+        self.best_epoch = 0
+        
+        # Metrics tracking
+        self.metrics_history = defaultdict(list)
+        
+        # Initialize trackio
+        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."""
+        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)
+            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()
+            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.no_grad()
+    def validate(self, epoch: int) -> dict:
+        """Validate the student model."""
+        self.student.eval()
+        self.teacher.eval()
+        
+        val_losses = defaultdict(float)
+        angular_errors = []
+        pitch_errors = []
+        yaw_errors = []
+        num_batches = 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."""
+        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()
+        
+        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
+                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
+            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
+
+
+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
+    
+    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')
+    os.makedirs(save_dir, exist_ok=True)
+    
+    for epoch in range(epochs):
+        # Training
+        teacher.train()
+        train_loss_total = 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()
+            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
+        teacher.eval()
+        val_loss_total = 0.0
+        val_angular = 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
+            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
+    device = torch.device(args.device if torch.cuda.is_available() else 'cpu')
+    print(f"Using device: {device}")
+    
+    # Create dataloaders
+    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",
+    )
+    
+    student = PriviGazeStudent()
+    
+    print(f"Teacher parameters: {count_parameters(teacher):,}")
+    print(f"Student parameters: {count_parameters(student):,}")
+    
+    # Pre-train teacher if needed
+    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
+    if args.teacher_path:
+        print(f"\nLoading teacher from: {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 = []
+        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
+        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
+
+
+if __name__ == "__main__":
+    main()