experiments/train_dann.py

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)