| 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 |
|
|
| |
| 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): |
| 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): |
| model.train() |
| running_loss = 0.0 |
| running_aw = 0.0 |
| running_domain = 0.0 |
| |
| alpha = get_dann_alpha(epoch, max_epochs) |
| |
| pbar = tqdm(loader, desc=f"Epoch {epoch} [DANN + AdaLN, alpha={alpha:.2f}]") |
| for batch in pbar: |
| patches, landmarks, gaze, t_p_logits, t_yaw_logits, domains = [b.to(device) for b in batch] |
| teacher_logits = (t_p_logits, t_yaw_logits) |
| |
| optimizer.zero_grad() |
| |
| |
| s_p_logits, s_y_logits, domain_logits = model(patches, landmarks, state='A', alpha=alpha, domain_id=domains) |
| |
| |
| 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) |
| |
| |
| d_loss = domain_criterion(domain_logits, domains) |
| |
| total_loss = gaze_loss + w_domain * d_loss |
| |
| 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_domain += d_loss.item() |
| |
| pbar.set_postfix({'loss': f"{total_loss.item():.4f}", 'aw': f"{l_aw:.4f}", 'dom': f"{d_loss.item():.4f}"}) |
| |
| return running_loss / len(loader), running_aw / len(loader), running_domain / len(loader) |
|
|
| def train_dann(h5_dir, source_suffix='_v16_new.h5', target_file='gaze360_robust_v16_new.h5', num_epochs=10): |
| device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') |
| print(f"EXPERIMENT 2: DANN + AdaLN (Clean Data) on {device}") |
| |
| all_files = os.listdir(h5_dir) |
| source_files = [os.path.join(h5_dir, f) for f in all_files if f.endswith(source_suffix) and f.startswith('p')] |
| target_path = os.path.join(h5_dir, target_file) |
| |
| if not source_files: |
| print(f"Waiting for source files with suffix {source_suffix}...") |
| return |
|
|
| 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=32, shuffle=True, num_workers=4) |
| |
| model = LIPEV2Student().to(device) |
| criterion = GazeDistillationLoss() |
| domain_criterion = torch.nn.CrossEntropyLoss() |
| optimizer = optim.Adam(model.parameters(), lr=1e-4) |
| |
| os.makedirs('checkpoints/dann_adaln', exist_ok=True) |
| |
| for epoch in range(1, num_epochs + 1): |
| w_aw, w_kd = 1.0, 0.5 |
| t_loss, l_aw, d_loss = train_dann_epoch(model, train_loader, criterion, domain_criterion, optimizer, device, epoch, num_epochs, w_aw, w_kd) |
| print(f"Epoch {epoch}: Total {t_loss:.4f}, AW {l_aw:.4f}, Dom {d_loss:.4f}") |
| |
| if epoch == num_epochs: |
| torch.save(model.state_dict(), f'checkpoints/dann_adaln/student_adaln_final.pt') |
|
|
| if __name__ == '__main__': |
| train_dann('data/processed', num_epochs=5) |
|
|
