experiments/train_stabilizer_v6_final.py

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)