Initial release of LIPE V2 GOLD

a10ba7f 5 months ago

9.52 kB

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)

	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)