src/training/trainer.py

"""
MultiSense-DF — Training Loop with Weighted Multi-Loss
Supports 3-stage training protocol
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.cuda.amp import autocast, GradScaler
from torch.optim import AdamW
from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR, SequentialLR
from sklearn.metrics import roc_auc_score
import numpy as np
from tqdm import tqdm


# ── Loss Functions ─────────────────────────────────────────────────────────────

class MultiModalLoss(nn.Module):
    """
    Weighted sum of:
      - global BCE        (weight 1.0)
      - visual BCE        (weight 0.3)
      - audio BCE         (weight 0.3)
      - lip-sync BCE      (weight 0.3)
    """

    def __init__(self, w_global=1.0, w_modal=0.3):
        super().__init__()
        self.w_global = w_global
        self.w_modal  = w_modal
        self.bce = nn.BCEWithLogitsLoss()

    def forward(self, outputs, labels):
        labels = labels.float().unsqueeze(1)

        loss_global  = self.bce(outputs['global_logit'], labels)
        loss_visual  = self.bce(outputs['per_mod_logits']['visual'],  labels)
        loss_audio   = self.bce(outputs['per_mod_logits']['audio'],   labels)
        loss_lipsync = self.bce(outputs['per_mod_logits']['lipsync'], labels)

        total = (self.w_global * loss_global +
                 self.w_modal  * loss_visual  +
                 self.w_modal  * loss_audio   +
                 self.w_modal  * loss_lipsync)

        return total, {
            'global':  loss_global.item(),
            'visual':  loss_visual.item(),
            'audio':   loss_audio.item(),
            'lipsync': loss_lipsync.item()
        }


# ── Trainer ────────────────────────────────────────────────────────────────────

class Trainer:
    def __init__(self, model, train_loader, val_loader,
                 lr=1e-4, weight_decay=0.01, epochs=30,
                 warmup_steps=2000, device='cuda',
                 save_path='checkpoints/multisense_df.pt',
                 use_amp=True):

        self.model        = model.to(device)
        self.train_loader = train_loader
        self.val_loader   = val_loader
        self.device       = device
        self.epochs       = epochs
        self.save_path    = save_path
        self.use_amp      = use_amp and (device == 'cuda')

        self.criterion = MultiModalLoss()
        self.optimizer = AdamW(
            filter(lambda p: p.requires_grad, model.parameters()),
            lr=lr, weight_decay=weight_decay
        )

        # Warmup then cosine annealing
        warmup = LinearLR(self.optimizer, start_factor=0.01,
                          end_factor=1.0, total_iters=warmup_steps)
        cosine = CosineAnnealingLR(self.optimizer, T_max=epochs - 10)
        self.scheduler = SequentialLR(
            self.optimizer, schedulers=[warmup, cosine],
            milestones=[warmup_steps]
        )

        self.scaler  = GradScaler(enabled=self.use_amp)
        self.best_auc = 0.0
        self.history  = {'train_loss': [], 'val_loss': [], 'val_auc': []}

    def train_epoch(self, epoch):
        self.model.train()
        total_loss = 0.0
        all_probs, all_labels = [], []

        pbar = tqdm(self.train_loader, desc=f'Epoch {epoch} [Train]', leave=False)
        for batch in pbar:
            frames      = batch['frames'].to(self.device)
            waveform    = batch['waveform'].to(self.device)
            mouth_crops = batch['mouth_crops'].to(self.device)
            mel_specs   = batch['mel_specs'].to(self.device)
            labels      = batch['label'].to(self.device)

            self.optimizer.zero_grad()
            with autocast(enabled=self.use_amp):
                outputs = self.model(frames, waveform, mouth_crops, mel_specs)
                loss, loss_dict = self.criterion(outputs, labels)

            self.scaler.scale(loss).backward()
            self.scaler.unscale_(self.optimizer)
            torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
            self.scaler.step(self.optimizer)
            self.scaler.update()
            self.scheduler.step()

            total_loss += loss.item()
            probs = torch.sigmoid(outputs['global_logit']).detach().cpu().numpy()
            all_probs.extend(probs.flatten())
            all_labels.extend(labels.cpu().numpy())
            pbar.set_postfix(loss=f'{loss.item():.4f}')

        avg_loss = total_loss / len(self.train_loader)
        auc = roc_auc_score(all_labels, all_probs) if len(set(all_labels)) > 1 else 0.0
        return avg_loss, auc

    @torch.no_grad()
    def val_epoch(self, epoch):
        self.model.eval()
        total_loss = 0.0
        all_probs, all_labels = [], []

        for batch in tqdm(self.val_loader, desc=f'Epoch {epoch} [Val]', leave=False):
            frames      = batch['frames'].to(self.device)
            waveform    = batch['waveform'].to(self.device)
            mouth_crops = batch['mouth_crops'].to(self.device)
            mel_specs   = batch['mel_specs'].to(self.device)
            labels      = batch['label'].to(self.device)

            with autocast(enabled=self.use_amp):
                outputs = self.model(frames, waveform, mouth_crops, mel_specs)
                loss, _ = self.criterion(outputs, labels)

            total_loss += loss.item()
            probs = torch.sigmoid(outputs['global_logit']).cpu().numpy()
            all_probs.extend(probs.flatten())
            all_labels.extend(labels.cpu().numpy())

        avg_loss = total_loss / len(self.val_loader)
        auc = roc_auc_score(all_labels, all_probs) if len(set(all_labels)) > 1 else 0.0
        return avg_loss, auc

    def fit(self):
        print(f'Training on {self.device} | AMP={self.use_amp}')
        for epoch in range(1, self.epochs + 1):
            train_loss, train_auc = self.train_epoch(epoch)
            val_loss,   val_auc   = self.val_epoch(epoch)

            self.history['train_loss'].append(train_loss)
            self.history['val_loss'].append(val_loss)
            self.history['val_auc'].append(val_auc)

            print(f'Epoch {epoch:03d} | '
                  f'Train Loss={train_loss:.4f} AUC={train_auc:.4f} | '
                  f'Val Loss={val_loss:.4f} AUC={val_auc:.4f}')

            if val_auc > self.best_auc:
                self.best_auc = val_auc
                torch.save({
                    'epoch': epoch,
                    'model_state': self.model.state_dict(),
                    'optimizer_state': self.optimizer.state_dict(),
                    'best_auc': self.best_auc,
                    'history': self.history
                }, self.save_path)
                print(f'  ✓ Saved best model (AUC={self.best_auc:.4f})')

        print(f'\nTraining complete. Best Val AUC: {self.best_auc:.4f}')
        return self.history