import torch import torch.nn as nn import torch.nn.functional as F import torch.optim as optim from torch.utils.data import DataLoader, ConcatDataset import os import time import numpy as np from tqdm import tqdm import sys from pathlib import Path # Disable HDF5 file locking for Windows compatibility os.environ["HDF5_USE_FILE_LOCKING"] = "FALSE" # 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_dann_alpha(epoch, max_epochs): """ Alpha schedule for DANN: starts from 0 and grows to 1. """ p = float(epoch) / max_epochs alpha = 2. / (1. + np.exp(-10 * p)) - 1 return alpha def train_dann_epoch(model, loader, criterion, domain_criterion, optimizer, device, epoch, max_epochs, w_aw, w_kd, w_domain=0.1, w_consist=0.5): model.train() running_loss = 0.0 running_aw = 0.0 running_kd = 0.0 running_domain = 0.0 running_consist = 0.0 alpha = get_dann_alpha(epoch, max_epochs) pbar = tqdm(loader, desc=f"Epoch {epoch} [alpha={alpha:.2f}]") for batch in pbar: # GazeDataset now returns (patch, landmarks, gaze, t_pitch, t_yaw, domain_id) patches, landmarks, gaze, t_p_logits, t_yaw_logits, domains = [b.to(device) for b in batch] # Check if we have valid teacher labels if torch.abs(t_p_logits).sum() > 0: teacher_logits = (t_p_logits, t_yaw_logits) else: teacher_logits = None optimizer.zero_grad() # 1. Standard Forward Pass s_p_logits, s_y_logits, domain_logits = model(patches, landmarks, state='A', alpha=alpha, domain_id=domains) # 2. Gaze Loss criterion.w_aw = w_aw criterion.w_kd = w_kd gaze_loss, l_aw, l_kd = criterion((s_p_logits, s_y_logits), gaze, teacher_logits) # 3. Domain Loss d_loss = domain_criterion(domain_logits, domains) # 4. Flip Consistency Loss patches_flipped = torch.flip(patches, dims=[3]) s_p_flip, s_y_flip, _ = model(patches_flipped, landmarks, state='A', alpha=alpha, domain_id=domains) def get_deg(logits): idx = torch.arange(90).float().to(device) prob = torch.softmax(logits, dim=1) return torch.sum(prob * idx, dim=1) * 2 - 90 p_deg = get_deg(s_p_logits) y_deg = get_deg(s_y_logits) p_deg_f = get_deg(s_p_flip) y_deg_f = get_deg(s_y_flip) loss_consist = F.mse_loss(p_deg, p_deg_f) + F.mse_loss(y_deg, -y_deg_f) total_loss = gaze_loss + w_domain * d_loss + w_consist * loss_consist total_loss.backward() optimizer.step() running_loss += total_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 running_domain += d_loss.item() running_consist += loss_consist.item() pbar.set_postfix({'loss': f"{total_loss.item():.4f}", 'aw': f"{l_aw:.4f}", 'con': f"{loss_consist.item():.4f}"}) return running_loss / len(loader), running_aw / len(loader), running_kd / len(loader), running_domain / len(loader), running_consist / len(loader) def train_dann(h5_dir, target_file='gaze360_robust_v16_train_A.h5', num_epochs=100, batch_size=32, baseline_model='checkpoints/baseline_v16/best_student_p11.pt', test_participant=None): device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') print(f"DANN + AdaLN + Consistency Training on {device}") all_files = os.listdir(h5_dir) source_files = [] for f in all_files: if f.endswith('_v16_new.h5') and f.startswith('p'): if test_participant and f.startswith(test_participant): print(f"LOPO: Skipping {f} for training.") continue f_path = os.path.join(h5_dir, f) if os.path.getsize(f_path) > 1024 * 1024: source_files.append(f_path) target_path = os.path.join(h5_dir, target_file) if not source_files: print("No robust source files found! Using available _v16.h5 files as fallback.") for f in all_files: if f.endswith('_v16.h5') and f.startswith('p'): if test_participant and f.startswith(test_participant): continue f_path = os.path.join(h5_dir, f) if os.path.getsize(f_path) > 1024 * 1024: source_files.append(f_path) print(f"Source files: {len(source_files)}") print(f"Target file: {target_path}") num_cpus = os.cpu_count() or 2 source_ds = GazeDataset(source_files, transform=True, domain_id=0) target_ds = GazeDataset([target_path], transform=True, domain_id=1) train_ds = ConcatDataset([source_ds, target_ds]) train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True, num_workers=min(num_cpus, 4)) model = LIPEV2Student().to(device) if baseline_model and os.path.exists(baseline_model): print(f"Initializing with baseline weights: {baseline_model}") state_dict = torch.load(baseline_model, map_location=device) model.load_state_dict(state_dict, strict=False) criterion = GazeDistillationLoss() domain_criterion = torch.nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=5e-5) out_dir = f'checkpoints/dann_lopo/{test_participant}' if test_participant else 'checkpoints/dann' os.makedirs(out_dir, exist_ok=True) for epoch in range(1, num_epochs + 1): metrics = train_dann_epoch(model, train_loader, criterion, domain_criterion, optimizer, device, epoch, num_epochs, 1.0, 0.5) t_loss, l_aw, l_kd, d_loss, c_loss = metrics print(f"Epoch {epoch}: Total {t_loss:.4f}, AW {l_aw:.4f}, KD {l_kd:.4f}, Dom {d_loss:.4f}, Consist {c_loss:.4f}") if epoch % 20 == 0 or epoch == num_epochs: save_path = os.path.join(out_dir, f'student_e{epoch}.pt') torch.save(model.state_dict(), save_path) if __name__ == '__main__': import argparse parser = argparse.ArgumentParser() parser.add_argument('--test_participant', type=str, default=None) parser.add_argument('--num_epochs', type=int, default=100) parser.add_argument('--batch_size', type=int, default=32) parser.add_argument('--h5_dir', type=str, default='data/processed') args = parser.parse_args() train_dann(args.h5_dir, num_epochs=args.num_epochs, batch_size=args.batch_size, test_participant=args.test_participant)