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)