""" 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