import torch import torch.nn as nn 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.data.dataset import GazeDataset # --- 1. Smooth Adaptive Wing Loss --- class SmoothAWLoss(nn.Module): def __init__(self, omega=8.0, alpha=1.5, theta=0.5, epsilon=1.0): super(SmoothAWLoss, self).__init__() self.omega = omega self.alpha = alpha self.theta = theta self.epsilon = epsilon def forward(self, y_pred, y_true): delta_y = (y_true - y_pred).abs() device = y_pred.device mask = delta_y < self.theta loss = torch.zeros_like(delta_y) loss[mask] = self.omega * torch.log(1 + torch.pow(delta_y[mask] / self.epsilon, self.alpha)) theta_eps = torch.tensor(self.theta / self.epsilon, device=device) A = self.omega * (1.0 / (1.0 + torch.pow(theta_eps, self.alpha))) * \ (self.alpha * torch.pow(theta_eps, self.alpha - 1.0) * (1.0 / self.epsilon)) B = A * self.theta - self.omega * torch.log(1.0 + torch.pow(theta_eps, self.alpha)) loss[~mask] = A * delta_y[~mask] - B return loss.mean() # --- 2. V6 Model: DualPool + Shuffle + LayerNorm --- class DualPoolMiniConv(nn.Module): def __init__(self): super(DualPoolMiniConv, self).__init__() self.conv = nn.Sequential( nn.Conv2d(1, 16, kernel_size=3, padding=0), nn.ReLU(inplace=True), nn.Conv2d(16, 32, kernel_size=3, padding=0), nn.ReLU(inplace=True), nn.Conv2d(32, 64, kernel_size=3, padding=0), nn.ReLU(inplace=True) ) self.avg_pool = nn.AdaptiveAvgPool2d(1) self.max_pool = nn.AdaptiveMaxPool2d(1) def forward(self, x): x = self.conv(x) return torch.cat([self.avg_pool(x), self.max_pool(x)], dim=1).flatten(1) class TokenShuffle(nn.Module): def forward(self, x): batch_size = x.shape[0] # (Batch, 4 patches * 128 features) -> (Batch, 4, 128) x = x.view(batch_size, 4, 128) # Transpose to interleave: (Batch, 128, 4) x = x.transpose(1, 2).contiguous() return x.view(batch_size, -1) class LIPEV2StudentV6(nn.Module): def __init__(self): super(LIPEV2StudentV6, self).__init__() self.app_net = DualPoolMiniConv() self.shuffle = TokenShuffle() self.geo_net = nn.Sequential( nn.Linear(956, 256), nn.LayerNorm(256), # Individual sample stability nn.ReLU(inplace=True), nn.Linear(256, 256), nn.ReLU(inplace=True) ) # Plan V6: LayerNorm for Cross-Domain Stability self.post_concat_ln = nn.LayerNorm(512 + 256) self.fusion = nn.Sequential( nn.Linear(512 + 256, 256), nn.ReLU(inplace=True), nn.Dropout(0.1), nn.Linear(256, 128), nn.ReLU(inplace=True) ) self.pitch_head = nn.Linear(128, 90) self.yaw_head = nn.Linear(128, 90) def forward(self, patches, landmarks): batch_size = patches.shape[0] p_h, p_w = patches.shape[2], patches.shape[3] app_feat = self.app_net(patches.view(-1, 1, p_h, p_w)).view(batch_size, -1) app_feat = self.shuffle(app_feat) geo_feat = self.geo_net(landmarks) combined = torch.cat([app_feat, geo_feat], dim=1) combined = self.post_concat_ln(combined) fused = self.fusion(combined) return self.pitch_head(fused), self.yaw_head(fused) # --- 3. Final Training Loop with OneCycleLR --- def train(h5_dir, test_participant='p11', num_epochs=75, max_lr=1.2e-4): device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') print(f"V6 THE STABILIZER | Participant: {test_participant} | Device: {device}") train_files = [os.path.join(h5_dir, f) for f in os.listdir(h5_dir) if f.endswith('_v16.h5') and not f.startswith(test_participant)] val_files = [os.path.join(h5_dir, f) for f in os.listdir(h5_dir) if f.startswith(test_participant) and f.endswith('_v16.h5')] # num_workers=3 to allow parallel training of 2 participants on 12-core CPU train_loader = DataLoader(GazeDataset(train_files, transform=True), batch_size=32, shuffle=True, num_workers=3) val_loader = DataLoader(GazeDataset(val_files, transform=False), batch_size=32, shuffle=False, num_workers=3) model = LIPEV2StudentV6().to(device) criterion = SmoothAWLoss() optimizer = optim.AdamW(model.parameters(), lr=max_lr/10, weight_decay=1e-2) # Plan V6: OneCycleLR Scheduler scheduler = optim.lr_scheduler.OneCycleLR( optimizer, max_lr=max_lr, steps_per_epoch=len(train_loader), epochs=num_epochs, pct_start=0.2, # 20% warmup anneal_strategy='cos' ) best_mae = float('inf') idx_tensor = torch.arange(90).float().to(device) for epoch in range(1, num_epochs + 1): model.train() running_loss = 0.0 pbar = tqdm(train_loader, desc=f"Epoch {epoch} Stabilizer ({test_participant})") for batch in pbar: batch = [b.to(device) for b in batch] optimizer.zero_grad() p_logits, y_logits = model(batch[0], batch[1]) s_p = torch.sum(torch.softmax(p_logits, dim=1) * idx_tensor, dim=1) * 2 - 90 s_y = torch.sum(torch.softmax(y_logits, dim=1) * idx_tensor, dim=1) * 2 - 90 loss = criterion(torch.stack([s_p, s_y], dim=1), batch[2] * (180.0 / np.pi)) loss.backward() optimizer.step() scheduler.step() # Step per batch running_loss += loss.item() pbar.set_postfix({'lr': f"{scheduler.get_last_lr()[0]:.2e}"}) model.eval() total_error, count = 0.0, 0 with torch.no_grad(): for batch in val_loader: batch = [b.to(device) for b in batch] p_logits, y_logits = model(batch[0], batch[1]) p_deg = (torch.sum(torch.softmax(p_logits, dim=1) * idx_tensor, dim=1) * 2 - 90) y_deg = (torch.sum(torch.softmax(y_logits, dim=1) * idx_tensor, dim=1) * 2 - 90) gt_deg = batch[2] * (180.0 / np.pi) total_error += (torch.abs(p_deg - gt_deg[:,0]) + torch.abs(y_deg - gt_deg[:,1])).sum().item() count += batch[2].shape[0] val_mae = total_error / (count * 2) print(f"Epoch {epoch}: Participant {test_participant}, Loss {running_loss/len(train_loader):.4f}, Val MAE {val_mae:.4f}") if val_mae < best_mae: best_mae = val_mae os.makedirs('checkpoints/stabilizer_v6', exist_ok=True) torch.save(model.state_dict(), f'checkpoints/stabilizer_v6/best_v6_{test_participant}.pt') with open('logs/best_epochs.log', 'a') as f: f.write(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] [V6-Final] Participant: {test_participant}, Epoch: {epoch}, MAE: {val_mae:.4f}\n") print(f"--- New Best ({test_participant}): {best_mae:.4f} ---") if __name__ == '__main__': import argparse parser = argparse.ArgumentParser() parser.add_argument('--participant', type=str, default='p11') args = parser.parse_args() train('data/processed', test_participant=args.participant, num_epochs=75, max_lr=1.2e-4)