train.py

"""
PriviGaze Training Script - Privileged Distillation for Gaze Estimation
"""
import os, sys, argparse, time
from pathlib import Path
from collections import defaultdict
import torch
from torch.optim import AdamW
from torch.optim.lr_scheduler import CosineAnnealingLR
import numpy as np

sys.path.insert(0, str(Path(__file__).parent))
from models.teacher import PriviGazeTeacher
from models.student import PriviGazeStudent, count_parameters
from models.distillation_loss import PriviGazeDistillationLoss, L2CSLoss, AngularLoss
from models.dataset import create_dataloaders

try:
    import trackio; HAS_TRACKIO = True
except ImportError:
    HAS_TRACKIO = False; print("Warning: trackio not installed.")


class DistillationTrainer:
    def __init__(self, teacher, student, dist_loss, train_loader, val_loader,
                 device, lr=1e-4, wd=1e-4, epochs=100, tproj="privi-gaze", trun="distill"):
        self.teacher = teacher.to(device)
        self.student = student.to(device)
        self.dist_loss = dist_loss.to(device)
        self.train_loader = train_loader
        self.val_loader = val_loader
        self.device = device
        self.epochs = epochs
        for p in self.teacher.parameters(): p.requires_grad = False
        self.teacher.eval()
        self.opt = AdamW(self.student.parameters(), lr=lr, weight_decay=wd)
        self.sched = CosineAnnealingLR(self.opt, T_max=epochs, eta_min=lr*0.01)
        self.best_val = float('inf')
        self.best_epoch = 0
        self.metrics = defaultdict(list)
        if HAS_TRACKIO:
            trackio.init(project=tproj, run_name=trun,
                config={'student_params': count_parameters(student),
                        'teacher_params': count_parameters(teacher), 'lr': lr, 'epochs': epochs})

    def train_epoch(self, epoch):
        self.student.train()
        losses = defaultdict(float)
        n = 0
        for bi, batch in enumerate(self.train_loader):
            le = batch['left_eye'].to(self.device)
            re = batch['right_eye'].to(self.device)
            fb = batch['face_blurred_gray'].to(self.device)
            fg = batch['face_gray'].to(self.device)
            pt = batch['pitch'].to(self.device)
            yt = batch['yaw'].to(self.device)
            with torch.no_grad():
                tp, ty, tplog, tylog, tf = self.teacher(le, re, fb)
            sp, sy, sf = self.student(fg)
            splog = self.student.pitch_head(sf)
            sylog = self.student.yaw_head(sf)
            loss, ld = self.dist_loss(sp, sy, splog, sylog, sf,
                                       tp, ty, tplog, tylog, tf, pt, yt)
            self.opt.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(self.student.parameters(), 1.0)
            self.opt.step()
            for k, v in ld.items(): losses[k] += v
            n += 1
            if bi % 100 == 0:
                print(f"Epoch {epoch} | Batch {bi} | " + " | ".join(f"{k}={v:.4f}" for k, v in ld.items()))
                if HAS_TRACKIO:
                    for k2, v2 in ld.items(): trackio.log({f"train/{k2}": v2})
        return {k: v/n for k, v in losses.items()}

    @torch.no_grad()
    def validate(self, epoch):
        self.student.eval()
        self.teacher.eval()
        losses = defaultdict(float)
        ae, pe, ye = [], [], []
        n = 0
        for batch in self.val_loader:
            le = batch['left_eye'].to(self.device)
            re = batch['right_eye'].to(self.device)
            fb = batch['face_blurred_gray'].to(self.device)
            fg = batch['face_gray'].to(self.device)
            pt = batch['pitch'].to(self.device)
            yt = batch['yaw'].to(self.device)
            tp, ty, tplog, tylog, tf = self.teacher(le, re, fb)
            sp, sy, sf = self.student(fg)
            splog = self.student.pitch_head(sf)
            sylog = self.student.yaw_head(sf)
            loss, ld = self.dist_loss(sp, sy, splog, sylog, sf,
                                       tp, ty, tplog, tylog, tf, pt, yt)
            for k, v in ld.items(): losses[k] += v
            n += 1
            aerr = torch.sqrt((sp-pt)**2 + (sy-yt)**2)
            ae.extend(aerr.cpu().tolist())
            pe.extend((sp-pt).abs().cpu().tolist())
            ye.extend((sy-yt).abs().cpu().tolist())
        for k in losses: losses[k] /= n
        losses['angular_mean'] = np.mean(ae)
        losses['angular_std'] = np.std(ae)
        losses['pitch_mean'] = np.mean(pe)
        losses['yaw_mean'] = np.mean(ye)
        return dict(losses)

    def train(self, save_dir="./checkpoints"):
        os.makedirs(save_dir, exist_ok=True)
        print(f"Distillation: {self.epochs} epochs | Student: {count_parameters(self.student):,} params")
        t0 = time.time()
        for epoch in range(self.epochs):
            te = time.time()
            tl = self.train_epoch(epoch)
            vl = self.validate(epoch)
            self.sched.step()
            lr = self.opt.param_groups[0]['lr']
            print(f"\n{'='*60}")
            print(f"Epoch {epoch}: train={tl.get('loss_total',0):.4f} val={vl.get('loss_total',0):.4f} angular={vl.get('angular_mean',0):.2f}deg")
            print(f"{'='*60}\n")
            for k, v in tl.items(): self.metrics[f'train_{k}'].append(v)
            for k, v in vl.items(): self.metrics[f'val_{k}'].append(v)
            vt = vl.get('loss_total', vl.get('angular_mean', float('inf')))
            if vt < self.best_val:
                self.best_val = vt
                self.best_epoch = epoch
                torch.save({'epoch': epoch, 'student_state_dict': self.student.state_dict(),
                            'opt_state_dict': self.opt.state_dict(), 'best_val': self.best_val,
                            'metrics': dict(self.metrics)}, os.path.join(save_dir, 'student_best.pt'))
                if HAS_TRACKIO: trackio.alert("New Best", f"Val {vt:.4f} @ epoch {epoch}", level="INFO")
            if epoch % 10 == 0:
                torch.save({'epoch': epoch, 'student_state_dict': self.student.state_dict(),
                            'opt_state_dict': self.opt.state_dict()},
                           os.path.join(save_dir, f'student_epoch_{epoch}.pt'))
            print(f"Epoch {epoch} took {time.time()-te:.1f}s, LR={lr:.2e}")
        print(f"\nDone! Best val: {self.best_val:.4f} @ epoch {self.best_epoch}")
        return self.best_val


def pretrain_teacher(teacher, train_loader, val_loader, device, lr=1e-4, epochs=50, save_dir="./checkpoints"):
    teacher = teacher.to(device)
    opt = AdamW(teacher.parameters(), lr=lr, weight_decay=1e-4)
    sched = CosineAnnealingLR(opt, T_max=epochs, eta_min=lr*0.01)
    ploss = L2CSLoss(gaze_bins=90)
    yloss = L2CSLoss(gaze_bins=90)
    aloss = AngularLoss()
    best = float('inf')
    os.makedirs(save_dir, exist_ok=True)
    for epoch in range(epochs):
        teacher.train()
        tloss = 0.0
        for batch in train_loader:
            le = batch['left_eye'].to(device)
            re = batch['right_eye'].to(device)
            fb = batch['face_blurred_gray'].to(device)
            pt = batch['pitch'].to(device)
            yt = batch['yaw'].to(device)
            pp, yp, pl, yl, _ = teacher(le, re, fb)
            loss = ploss(pl, pp, pt) + yloss(yl, yp, yt) + aloss(pp, yp, pt, yt)
            opt.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(teacher.parameters(), 1.0)
            opt.step()
            tloss += loss.item()
        tloss /= len(train_loader)
        teacher.eval()
        vloss = 0.0
        va = 0.0
        with torch.no_grad():
            for batch in val_loader:
                le = batch['left_eye'].to(device)
                re = batch['right_eye'].to(device)
                fb = batch['face_blurred_gray'].to(device)
                pt = batch['pitch'].to(device)
                yt = batch['yaw'].to(device)
                pp, yp, pl, yl, _ = teacher(le, re, fb)
                vloss += (ploss(pl, pp, pt) + yloss(yl, yp, yt)).item()
                va += torch.sqrt((pp-pt)**2 + (yp-yt)**2).mean().item()
        vloss /= len(val_loader)
        va /= len(val_loader)
        sched.step()
        print(f"Teacher Epoch {epoch}: train={tloss:.4f} val={vloss:.4f} angular={va:.2f}deg")
        if vloss < best:
            best = vloss
            torch.save(teacher.state_dict(), os.path.join(save_dir, 'teacher_best.pt'))
    return os.path.join(save_dir, 'teacher_best.pt')


def main():
    p = argparse.ArgumentParser(description="PriviGaze Training")
    p.add_argument('--mode', type=str, default='distill', choices=['pretrain_teacher','distill','both'])
    p.add_argument('--teacher-path', type=str, default=None)
    p.add_argument('--batch-size', type=int, default=32)
    p.add_argument('--epochs', type=int, default=100)
    p.add_argument('--teacher-epochs', type=int, default=50)
    p.add_argument('--lr', type=float, default=1e-4)
    p.add_argument('--weight-decay', type=float, default=1e-4)
    p.add_argument('--num-train', type=int, default=40000)
    p.add_argument('--num-val', type=int, default=5000)
    p.add_argument('--save-dir', type=str, default='./checkpoints')
    p.add_argument('--device', type=str, default='cuda')
    p.add_argument('--trackio-project', type=str, default='privi-gaze')
    p.add_argument('--trackio-run', type=str, default='distill-run')
    p.add_argument('--push-to-hub', action='store_true')
    p.add_argument('--hub-model-id', type=str, default=None)
    p.add_argument('--alpha-contrastive', type=float, default=0.5)
    p.add_argument('--alpha-mmd', type=float, default=0.1)
    p.add_argument('--alpha-logit', type=float, default=0.5)
    args = p.parse_args()

    device = torch.device(args.device if torch.cuda.is_available() else 'cpu')
    print(f"Device: {device}")
    train_loader, val_loader, test_loader = create_dataloaders(
        num_train=args.num_train, num_val=args.num_val, batch_size=args.batch_size)

    teacher = PriviGazeTeacher()
    student = PriviGazeStudent()
    print(f"Teacher: {count_parameters(teacher):,} params")
    print(f"Student: {count_parameters(student):,} params")

    if args.mode in ['pretrain_teacher', 'both']:
        print("\n=== Phase 1: Teacher Pre-training ===")
        tp = pretrain_teacher(teacher, train_loader, val_loader, device,
                              lr=args.lr, epochs=args.teacher_epochs, save_dir=args.save_dir)
        args.teacher_path = tp

    if args.teacher_path:
        print(f"\nLoading teacher: {args.teacher_path}")
        teacher.load_state_dict(torch.load(args.teacher_path, map_location=device))

    if args.mode in ['distill', 'both']:
        print("\n=== Phase 2: Distillation ===")
        dloss = PriviGazeDistillationLoss(
            gaze_bins=90, teacher_feature_dim=256, student_feature_dim=128,
            alpha_contrastive=args.alpha_contrastive, alpha_mmd=args.alpha_mmd,
            alpha_logit=args.alpha_logit)
        trainer = DistillationTrainer(teacher, student, dloss, train_loader, val_loader,
            device, lr=args.lr, wd=args.weight_decay, epochs=args.epochs,
            tproj=args.trackio_project, trun=args.trackio_run)
        trainer.train(save_dir=args.save_dir)

        print("\n=== Test ===")
        student.eval().to(device)
        terr = []
        with torch.no_grad():
            for batch in test_loader:
                fg = batch['face_gray'].to(device)
                pt = batch['pitch'].to(device)
                yt = batch['yaw'].to(device)
                sp, sy, _ = student(fg)
                terr.extend(torch.sqrt((sp-pt)**2 + (sy-yt)**2).cpu().tolist())
        me = np.mean(terr); se = np.std(terr)
        print(f"Test Angular Error: {me:.2f}deg +- {se:.2f}deg")

        if args.push_to_hub and args.hub_model_id:
            from huggingface_hub import HfApi
            mp = os.path.join(args.save_dir, 'student_final.pt')
            torch.save({'student_state_dict': student.state_dict(),
                        'config': {'params': count_parameters(student), 'test_err': me}}, mp)
            HfApi().upload_file(path_or_fileobj=mp, path_in_repo="student_model.pt", repo_id=args.hub_model_id)
            print(f"Pushed to: https://huggingface.co/{args.hub_model_id}")

if __name__ == "__main__":
    main()