experiments/train.py

import torch
import torch.optim as optim
from torch.utils.data import DataLoader
import os
import time
import numpy as np
from tqdm import tqdm
import sys
from pathlib import Path

# Add project root to path
sys.path.append(str(Path(__file__).parent.parent))

from src.models.student import LIPEV2Student
from src.models.loss import GazeDistillationLoss
from src.data.dataset import GazeDataset

def get_loss_weights(epoch):
    """
    Dynamic Weighting Schedule from PHASE3_PLAN.md
    """
    if epoch <= 5: # Stage 0
        w_aw, w_kd = 1.0, 0.5
    elif epoch <= 30: # Stage 1
        w_aw, w_kd = 0.2, 1.0
    elif epoch <= 40: # Transition Phase (Linear Annealing)
        # Epoch 31: start transition
        # Epoch 40: end transition
        t = (epoch - 30) / 10.0
        w_aw = 0.2 + t * (1.0 - 0.2)
        w_kd = 1.0 - t * (1.0 - 0.1)
    else: # Stage 2
        w_aw, w_kd = 1.0, 0.1
        
    return w_aw, w_kd

def train_epoch(model, loader, criterion, optimizer, device, epoch, w_aw, w_kd):
    model.train()
    running_loss = 0.0
    running_aw = 0.0
    running_kd = 0.0
    
    pbar = tqdm(loader, desc=f"Epoch {epoch} [w_aw={w_aw:.2f}, w_kd={w_kd:.2f}]")
    for batch in pbar:
        # Move all tensors to device
        batch = [b.to(device) for b in batch]
        
        # GazeDataset returns:
        # 3 items: (patch, landmarks, gaze)
        # 5 items: (patch, landmarks, gaze, t_p, t_y)
        # 6 items: (patch, landmarks, gaze, t_p, t_y, domain_id)
        if len(batch) >= 5:
            patches, landmarks, gaze, t_p_logits, t_yaw_logits = batch[:5]
            teacher_logits = (t_p_logits, t_yaw_logits)
        else:
            patches, landmarks, gaze = batch[:3]
            teacher_logits = None
            
        optimizer.zero_grad()
        
        # We always train in State A (Full mode)
        student_outputs = model(patches, landmarks, state='A')
        
        # Set weights dynamically
        criterion.w_aw = w_aw
        criterion.w_kd = w_kd
        
        loss, l_aw, l_kd = criterion(student_outputs, gaze, teacher_logits)
        
        loss.backward()
        optimizer.step()
        
        running_loss += loss.item()
        running_aw += l_aw.item() if isinstance(l_aw, torch.Tensor) else l_aw
        running_kd += l_kd.item() if isinstance(l_kd, torch.Tensor) else l_kd
        
        pbar.set_postfix({'loss': f"{loss.item():.4f}", 'aw': f"{l_aw:.4f}", 'kd': f"{l_kd:.4f}"})
        
    return running_loss / len(loader), running_aw / len(loader), running_kd / len(loader)

def validate(model, loader, device):
    model.eval()
    total_error = 0.0
    count = 0
    
    with torch.no_grad():
        for batch in loader:
            batch = [b.to(device) for b in batch]
            patches, landmarks, gaze = batch[0], batch[1], batch[2]
            
            outputs = model(patches, landmarks, state='A')
            if len(outputs) == 3:
                p_logits, y_logits, _ = outputs
            else:
                p_logits, y_logits = outputs
            
            # Convert logits to angles (Degrees)
            idx = torch.arange(90).float().to(device)
            p_prob = torch.softmax(p_logits, dim=1)
            y_prob = torch.softmax(y_logits, dim=1)
            
            p_deg = (torch.sum(p_prob * idx, dim=1) * 2 - 90)
            y_deg = (torch.sum(y_prob * idx, dim=1) * 2 - 90)
            
            # GT Gaze is in radians, convert to degrees
            gt_deg = gaze * (180.0 / np.pi)
            
            # Simple L1 error for now (Pitch/Yaw MAE)
            error = torch.abs(p_deg - gt_deg[:, 0]) + torch.abs(y_deg - gt_deg[:, 1])
            total_error += error.sum().item()
            count += gaze.shape[0]
            
    if count == 0:
        return float('inf')
        
    return total_error / (count * 2) # Average MAE per axis

class EarlyStopping:
    def __init__(self, patience=10, min_delta=0.001, verbose=True):
        self.patience = patience
        self.min_delta = min_delta
        self.verbose = verbose
        self.counter = 0
        self.best_loss = None
        self.early_stop = False

    def __call__(self, val_loss):
        if self.best_loss is None:
            self.best_loss = val_loss
        elif val_loss > self.best_loss - self.min_delta:
            self.counter += 1
            if self.verbose:
                print(f"EarlyStopping counter: {self.counter} out of {self.patience}")
            if self.counter >= self.patience:
                self.early_stop = True
        else:
            self.best_loss = val_loss
            self.counter = 0

def train(h5_dir, test_participant='p00', num_epochs=100, batch_size=64, lr=1e-4, no_kd=False, patience=15, version_suffix="", num_workers=4):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"Training on {device} {'(No KD)' if no_kd else ''} | Version: {version_suffix if version_suffix else 'baseline'}")
    
    # LOPO: Split files
    all_files = os.listdir(h5_dir)
    if version_suffix:
        relevant_files = [f for f in all_files if f.endswith(f"{version_suffix}.h5")]
    else:
        relevant_files = [f for f in all_files if f.endswith('.h5') and '_v' not in f]
        
    train_files = [os.path.join(h5_dir, f) for f in relevant_files if not f.startswith(test_participant)]
    val_files = [os.path.join(h5_dir, f) for f in relevant_files if f.startswith(test_participant)]
    
    if not train_files:
        print(f"No training files found for {test_participant} with suffix '{version_suffix}'.")
        return
        
    train_ds = GazeDataset(train_files, transform=True) # Enable augmentations for training
    val_ds = GazeDataset(val_files, transform=False) # Disable for validation
        
    # Subject-Level Weighted Sampling (PHASE3_PLAN.md)
    # Ensure each participant contributes equally
    sample_weights = []
    # Count samples per file
    import h5py
    counts = {}
    for i, f_path in enumerate(train_files):
        with h5py.File(f_path, 'r') as f:
            counts[f_path] = f['left_patches'].shape[0]
            
    for i, j in train_ds.indices:
        f_path = train_ds.h5_files[i]
        sample_weights.append(1.0 / counts[f_path])
        
    sampler = torch.utils.data.WeightedRandomSampler(sample_weights, len(sample_weights))
    
    train_loader = DataLoader(train_ds, batch_size=batch_size, sampler=sampler, num_workers=num_workers, pin_memory=True)
    val_loader = DataLoader(val_ds, batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=True)

    
    model = LIPEV2Student().to(device)
    criterion = GazeDistillationLoss()
    optimizer = optim.Adam(model.parameters(), lr=lr)
    # Refinement: Phase 5 Slow and Steady Scheduler
    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=15)
    
    early_stopping = EarlyStopping(patience=patience, verbose=True)
    
    best_mae = float('inf')
    os.makedirs('checkpoints', exist_ok=True)
    os.makedirs('logs', exist_ok=True)
    
    for epoch in range(1, num_epochs + 1):
        if no_kd:
            w_aw, w_kd = 1.0, 0.0
        else:
            w_aw, w_kd = get_loss_weights(epoch)
        
        train_loss, train_aw, train_kd = train_epoch(model, train_loader, criterion, optimizer, device, epoch, w_aw, w_kd)
        val_mae = validate(model, val_loader, device)
        
        # Get current learning rate for logging
        current_lr = optimizer.param_groups[0]['lr']
        print(f"Epoch {epoch}: Train Loss {train_loss:.4f}, Val MAE {val_mae:.4f}, LR {current_lr:.6f}")
        
        if val_mae < best_mae:
            best_mae = val_mae
            torch.save(model.state_dict(), f'checkpoints/best_student_{test_participant}.pt')
            log_msg = f"New best model saved at epoch {epoch}! (MAE: {best_mae:.4f})"
            print(log_msg)
            
            # Log to file
            with open('logs/best_epochs.log', 'a', encoding='utf-8') as f:
                f.write(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] Participant: {test_participant}, Epoch: {epoch}, MAE: {best_mae:.4f}, Version: {version_suffix}\n")
            
        early_stopping(val_mae)
        if early_stopping.early_stop:
            print(f"Early stopping triggered at epoch {epoch}")
            break
            
        scheduler.step(val_mae)

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--h5_dir', type=str, default='data/processed')
    parser.add_argument('--test_participant', type=str, default='p00')
    parser.add_argument('--num_epochs', type=int, default=100)
    parser.add_argument('--batch_size', type=int, default=32)
    parser.add_argument('--lr', type=float, default=1e-4)
    parser.add_argument('--patience', type=int, default=15)
    parser.add_argument('--no_kd', action='store_true', help='Disable Knowledge Distillation')
    parser.add_argument('--version_suffix', type=str, default='', help='Suffix for h5 files (e.g. _v16)')
    parser.add_argument('--num_workers', type=int, default=4)
    args = parser.parse_args()
    
    train(args.h5_dir, 
          test_participant=args.test_participant, 
          num_epochs=args.num_epochs, 
          batch_size=args.batch_size, 
          lr=args.lr,
          no_kd=args.no_kd,
          patience=args.patience,
          version_suffix=args.version_suffix,
          num_workers=args.num_workers)