Add train_v2.py: OGM-GE training loop + asymmetric LRs + aux loss logging

src/train_v2.py

@@ -0,0 +1,585 @@
+"""
+Multimodal PC Fault Detection - Training Script v2
+====================================================
+Changes from v1:
+  - OGM-GE gradient modulation after loss.backward(), before optimizer.step()
+  - Asymmetric learning rates: higher for visual branch, lower for audio
+  - Auxiliary loss logging (loss_fusion, loss_visual, loss_audio per epoch)
+  - OGM-GE stats logging (visual_conf, audio_conf, modulation coefficients)
+  - Uses models_v2 (auxiliary heads) and dataset_real (real data)
+
+Usage:
+  python train_v2.py --mode multimodal --finetune lora --eval_robustness
+  python train_v2.py --mode visual_only --finetune lora --no_push
+  python train_v2.py --mode audio_only --finetune lora --no_push
+  python train_v2.py --mode multimodal --finetune full --lr 2e-5
+  python train_v2.py --quick_test --no_push
+
+References:
+  OGM-GE: Peng et al., "Balanced Multimodal Learning via On-the-fly Gradient 
+  Modulation", CVPR 2022
+"""
+
+import os, sys, json, argparse, time
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+from torch.optim import AdamW
+from torch.optim.lr_scheduler import OneCycleLR
+from sklearn.metrics import accuracy_score, f1_score, confusion_matrix, precision_recall_fscore_support
+
+from config import ExperimentConfig, FAULT_CLASSES, NUM_CLASSES
+from dataset_real import RealPCFaultDataset as PCFaultDataset, multimodal_collate_fn
+from models_v2 import create_model, get_processors, OGMGEModulator
+
+
+def compute_metrics(preds, labels, class_names=FAULT_CLASSES):
+    """Compute accuracy, F1, precision, recall, and confusion matrix."""
+    accuracy = accuracy_score(labels, preds)
+    precision, recall, f1, support = precision_recall_fscore_support(
+        labels, preds, average=None, labels=range(len(class_names)), zero_division=0)
+    macro_f1 = f1_score(labels, preds, average="macro", zero_division=0)
+    weighted_f1 = f1_score(labels, preds, average="weighted", zero_division=0)
+    conf_matrix = confusion_matrix(labels, preds, labels=range(len(class_names)))
+
+    metrics = {
+        "accuracy": accuracy,
+        "macro_f1": macro_f1,
+        "weighted_f1": weighted_f1,
+        "confusion_matrix": conf_matrix.tolist(),
+        "per_class": {},
+    }
+    for i, name in enumerate(class_names):
+        metrics["per_class"][name] = {
+            "precision": precision[i], "recall": recall[i],
+            "f1": f1[i], "support": int(support[i]),
+        }
+    return metrics
+
+
+class MultimodalTrainerV2:
+    """
+    Training loop v2 with OGM-GE gradient modulation.
+    
+    Key differences from v1:
+    1. Three separate parameter groups with asymmetric LRs:
+       - visual_branch: higher LR (visual_lr_multiplier × base_lr)
+       - audio_branch: lower LR (audio_lr_multiplier × base_lr)
+       - fusion + auxiliary heads: base LR
+    2. OGM-GE applied after backward(), before optimizer.step()
+    3. Logs auxiliary losses and OGM-GE stats per epoch
+    """
+
+    def __init__(self, model, train_dataset, val_dataset, config, device,
+                 use_ogm=True, ogm_alpha=0.3, ogm_noise_sigma=0.1,
+                 visual_lr_multiplier=3.0, audio_lr_multiplier=0.5):
+        self.model = model.to(device)
+        self.device = device
+        self.config = config
+        self.use_ogm = use_ogm and (model.mode == "multimodal")
+
+        # OGM-GE modulator
+        if self.use_ogm:
+            self.ogm = OGMGEModulator(alpha=ogm_alpha, noise_sigma=ogm_noise_sigma)
+            print(f"[Trainer v2] OGM-GE enabled: alpha={ogm_alpha}, noise_sigma={ogm_noise_sigma}")
+        else:
+            self.ogm = None
+
+        # Data loaders
+        self.train_loader = DataLoader(
+            train_dataset,
+            batch_size=config.per_device_train_batch_size,
+            shuffle=True,
+            collate_fn=multimodal_collate_fn,
+            num_workers=2,
+            pin_memory=True,
+            drop_last=True)
+        self.val_loader = DataLoader(
+            val_dataset,
+            batch_size=config.per_device_eval_batch_size,
+            shuffle=False,
+            collate_fn=multimodal_collate_fn,
+            num_workers=2,
+            pin_memory=True)
+
+        # Asymmetric parameter groups
+        param_groups = self._get_param_groups(visual_lr_multiplier, audio_lr_multiplier)
+        self.optimizer = AdamW(param_groups, weight_decay=config.weight_decay)
+
+        total_steps = (len(self.train_loader) * config.num_epochs
+                       // config.gradient_accumulation_steps)
+        self.scheduler = OneCycleLR(
+            self.optimizer,
+            max_lr=[pg["lr"] for pg in param_groups],
+            total_steps=max(total_steps, 1),
+            pct_start=config.warmup_ratio,
+            anneal_strategy="cos")
+
+        # Mixed precision
+        self.scaler = (torch.amp.GradScaler("cuda")
+                       if config.fp16 and device.type == "cuda" else None)
+
+        # Tracking
+        self.best_metric = 0.0
+        self.best_epoch = 0
+        self.history = {
+            "train_loss": [], "val_loss": [],
+            "val_accuracy": [], "val_macro_f1": [],
+            # v2 additions
+            "train_loss_fusion": [], "train_loss_visual": [], "train_loss_audio": [],
+            "ogm_visual_conf": [], "ogm_audio_conf": [],
+            "ogm_coeff_visual": [], "ogm_coeff_audio": [],
+        }
+
+    def _get_param_groups(self, visual_lr_multiplier, audio_lr_multiplier):
+        """
+        Create 3 parameter groups with asymmetric learning rates.
+        
+        For LoRA mode: uses lora_learning_rate as base.
+        Visual branch gets multiplier > 1 (boost weak modality).
+        Audio branch gets multiplier < 1 (slow down dominant modality).
+        Fusion head + auxiliary heads get base LR.
+        """
+        visual_params = []
+        audio_params = []
+        fusion_params = []  # fusion head + auxiliary heads
+
+        for name, param in self.model.named_parameters():
+            if not param.requires_grad:
+                continue
+            if "visual_branch" in name:
+                visual_params.append(param)
+            elif "audio_branch" in name:
+                audio_params.append(param)
+            else:
+                fusion_params.append(param)
+
+        # Determine base LR
+        base_lr = self.config.lora_learning_rate  # default: 5e-3
+
+        groups = []
+        if visual_params:
+            vlr = base_lr * visual_lr_multiplier
+            groups.append({"params": visual_params, "lr": vlr, "name": "visual_branch"})
+            print(f"[Trainer v2] visual_branch: {len(visual_params)} tensors, lr={vlr:.2e}")
+        if audio_params:
+            alr = base_lr * audio_lr_multiplier
+            groups.append({"params": audio_params, "lr": alr, "name": "audio_branch"})
+            print(f"[Trainer v2] audio_branch: {len(audio_params)} tensors, lr={alr:.2e}")
+        if fusion_params:
+            groups.append({"params": fusion_params, "lr": base_lr, "name": "fusion_heads"})
+            print(f"[Trainer v2] fusion_heads: {len(fusion_params)} tensors, lr={base_lr:.2e}")
+
+        if not groups:
+            raise ValueError("No trainable parameters!")
+        return groups
+
+    def train_epoch(self, epoch):
+        """Train one epoch with OGM-GE gradient modulation."""
+        self.model.train()
+        total_loss = 0.0
+        total_loss_fusion = 0.0
+        total_loss_visual = 0.0
+        total_loss_audio = 0.0
+        num_batches = 0
+
+        # OGM-GE stats accumulators
+        ogm_v_confs, ogm_a_confs = [], []
+        ogm_cv, ogm_ca = [], []
+
+        self.optimizer.zero_grad()
+
+        for batch_idx, batch in enumerate(self.train_loader):
+            pv = batch["pixel_values"].to(self.device)
+            av = batch["audio_values"].to(self.device)
+            labels = batch["labels"].to(self.device)
+
+            # Forward pass
+            if self.scaler:
+                with torch.amp.autocast("cuda"):
+                    outputs = self.model(pixel_values=pv, audio_values=av, labels=labels)
+                    loss = outputs["loss"] / self.config.gradient_accumulation_steps
+                self.scaler.scale(loss).backward()
+            else:
+                outputs = self.model(pixel_values=pv, audio_values=av, labels=labels)
+                loss = outputs["loss"] / self.config.gradient_accumulation_steps
+                loss.backward()
+
+            # Accumulate losses
+            total_loss += loss.item() * self.config.gradient_accumulation_steps
+            num_batches += 1
+            if "loss_fusion" in outputs:
+                total_loss_fusion += outputs["loss_fusion"]
+                total_loss_visual += outputs["loss_visual"]
+                total_loss_audio += outputs["loss_audio"]
+
+            # Collect OGM-GE stats every batch (but only apply at accumulation boundary)
+            if (self.use_ogm and self.ogm is not None
+                    and "visual_logits" in outputs and "audio_logits" in outputs):
+                _cv, _ca, _stats = self.ogm.compute_modulation_coefficients(
+                    outputs["visual_logits"], outputs["audio_logits"], labels)
+                ogm_v_confs.append(_stats["visual_conf"])
+                ogm_a_confs.append(_stats["audio_conf"])
+                ogm_cv.append(_stats["coeff_visual"])
+                ogm_ca.append(_stats["coeff_audio"])
+
+            # Optimizer step (at accumulation boundary)
+            if (batch_idx + 1) % self.config.gradient_accumulation_steps == 0:
+                if self.scaler:
+                    self.scaler.unscale_(self.optimizer)
+
+                # ==== OGM-GE: modulate gradients AFTER unscale, BEFORE step ====
+                if (self.use_ogm and self.ogm is not None and ogm_cv):
+                    # Use the most recent coefficients
+                    self.ogm.apply_gradient_modulation(
+                        self.model, ogm_cv[-1], ogm_ca[-1])
+
+                torch.nn.utils.clip_grad_norm_(
+                    self.model.parameters(), self.config.max_grad_norm)
+
+                if self.scaler:
+                    self.scaler.step(self.optimizer)
+                    self.scaler.update()
+                else:
+                    self.optimizer.step()
+                self.scheduler.step()
+                self.optimizer.zero_grad()
+
+            # Logging
+            if (batch_idx + 1) % self.config.logging_steps == 0 or batch_idx == 0:
+                avg_loss = total_loss / num_batches
+                msg = (f"  [Epoch {epoch+1}] Step {batch_idx+1}/{len(self.train_loader)} "
+                       f"| Loss: {avg_loss:.4f} "
+                       f"| LR_v: {self.optimizer.param_groups[0]['lr']:.2e}")
+                if "loss_fusion" in outputs:
+                    msg += (f" | L_fus: {total_loss_fusion/num_batches:.4f}"
+                            f" L_vis: {total_loss_visual/num_batches:.4f}"
+                            f" L_aud: {total_loss_audio/num_batches:.4f}")
+                if ogm_cv:
+                    msg += (f" | OGM c_v: {ogm_cv[-1]:.3f}"
+                            f" c_a: {ogm_ca[-1]:.3f}")
+                print(msg)
+
+        # Epoch-level OGM stats
+        n = max(num_batches, 1)
+        epoch_stats = {
+            "train_loss": total_loss / n,
+            "loss_fusion": total_loss_fusion / n,
+            "loss_visual": total_loss_visual / n,
+            "loss_audio": total_loss_audio / n,
+        }
+        if ogm_v_confs:
+            epoch_stats["ogm_visual_conf"] = np.mean(ogm_v_confs)
+            epoch_stats["ogm_audio_conf"] = np.mean(ogm_a_confs)
+            epoch_stats["ogm_coeff_visual"] = np.mean(ogm_cv)
+            epoch_stats["ogm_coeff_audio"] = np.mean(ogm_ca)
+
+        return epoch_stats
+
+    @torch.no_grad()
+    def evaluate(self, modality_mask=None):
+        """Evaluate on validation set. Optionally mask a modality for robustness test."""
+        self.model.eval()
+        all_preds, all_labels = [], []
+        total_loss = 0.0
+        num_batches = 0
+
+        for batch in self.val_loader:
+            pv = batch["pixel_values"].to(self.device)
+            av = batch["audio_values"].to(self.device)
+            labels = batch["labels"].to(self.device)
+
+            if modality_mask:
+                if modality_mask.get("visual", 1.0) == 0.0:
+                    pv = torch.zeros_like(pv)
+                if modality_mask.get("audio", 1.0) == 0.0:
+                    av = torch.zeros_like(av)
+
+            outputs = self.model(pixel_values=pv, audio_values=av, labels=labels)
+            total_loss += outputs["loss"].item()
+            num_batches += 1
+            all_preds.extend(outputs["logits"].argmax(dim=-1).cpu().numpy())
+            all_labels.extend(labels.cpu().numpy())
+
+        metrics = compute_metrics(np.array(all_preds), np.array(all_labels))
+        metrics["val_loss"] = total_loss / max(num_batches, 1)
+        return metrics
+
+    def train(self):
+        """Full training loop with OGM-GE and detailed logging."""
+        print(f"\n{'='*60}")
+        print(f"Training v2: mode={self.model.mode}, epochs={self.config.num_epochs}, "
+              f"batch={self.config.per_device_train_batch_size}, device={self.device}")
+        print(f"OGM-GE: {'ENABLED' if self.use_ogm else 'DISABLED'}")
+        if self.model.mode == "multimodal":
+            print(f"Auxiliary loss weights: λ_visual={self.model.lambda_visual}, "
+                  f"λ_audio={self.model.lambda_audio}")
+        print(f"{'='*60}\n")
+
+        for epoch in range(self.config.num_epochs):
+            t0 = time.time()
+            train_stats = self.train_epoch(epoch)
+            val_metrics = self.evaluate()
+
+            # Print epoch summary
+            elapsed = time.time() - t0
+            print(f"\n[Epoch {epoch+1}/{self.config.num_epochs}] ({elapsed:.1f}s)")
+            loss_msg = f"  Train Loss: {train_stats['train_loss']:.4f}"
+            if train_stats.get("loss_fusion", 0) > 0:
+                loss_msg += (f" (fusion={train_stats['loss_fusion']:.4f} "
+                             f"visual={train_stats['loss_visual']:.4f} "
+                             f"audio={train_stats['loss_audio']:.4f})")
+            print(loss_msg)
+            print(f"  Val Loss: {val_metrics['val_loss']:.4f} "
+                  f"| Acc: {val_metrics['accuracy']:.4f} "
+                  f"| F1: {val_metrics['macro_f1']:.4f}")
+
+            if "ogm_visual_conf" in train_stats:
+                print(f"  OGM-GE: visual_conf={train_stats['ogm_visual_conf']:.4f} "
+                      f"audio_conf={train_stats['ogm_audio_conf']:.4f} "
+                      f"| coeff_v={train_stats['ogm_coeff_visual']:.4f} "
+                      f"coeff_a={train_stats['ogm_coeff_audio']:.4f}")
+
+            # Update history
+            self.history["train_loss"].append(train_stats["train_loss"])
+            self.history["val_loss"].append(val_metrics["val_loss"])
+            self.history["val_accuracy"].append(val_metrics["accuracy"])
+            self.history["val_macro_f1"].append(val_metrics["macro_f1"])
+            self.history["train_loss_fusion"].append(train_stats["loss_fusion"])
+            self.history["train_loss_visual"].append(train_stats["loss_visual"])
+            self.history["train_loss_audio"].append(train_stats["loss_audio"])
+            if "ogm_visual_conf" in train_stats:
+                self.history["ogm_visual_conf"].append(train_stats["ogm_visual_conf"])
+                self.history["ogm_audio_conf"].append(train_stats["ogm_audio_conf"])
+                self.history["ogm_coeff_visual"].append(train_stats["ogm_coeff_visual"])
+                self.history["ogm_coeff_audio"].append(train_stats["ogm_coeff_audio"])
+
+            # Save best model
+            if val_metrics[self.config.metric_for_best_model] > self.best_metric:
+                self.best_metric = val_metrics[self.config.metric_for_best_model]
+                self.best_epoch = epoch + 1
+                os.makedirs(self.config.output_dir, exist_ok=True)
+                torch.save({
+                    "model_state_dict": self.model.state_dict(),
+                    "epoch": epoch + 1,
+                    "metrics": val_metrics,
+                }, os.path.join(self.config.output_dir, "best_model.pt"))
+                print(f"  ✓ Best model saved (F1={self.best_metric:.4f})")
+
+        print(f"\nTraining complete. Best epoch={self.best_epoch}, "
+              f"Best F1={self.best_metric:.4f}")
+        return self.history
+
+    def run_robustness_evaluation(self):
+        """Test with missing modalities to evaluate robustness."""
+        print("\n=== Missing Modality Robustness Evaluation ===")
+        results = {}
+        scenarios = [
+            ("both_modalities", None),
+            ("visual_only", {"visual": 1.0, "audio": 0.0}),
+            ("audio_only", {"visual": 0.0, "audio": 1.0}),
+        ]
+        for name, mask in scenarios:
+            m = self.evaluate(modality_mask=mask)
+            results[name] = {"accuracy": m["accuracy"], "macro_f1": m["macro_f1"]}
+            print(f"  {name:20s}: Acc={m['accuracy']:.4f} F1={m['macro_f1']:.4f}")
+
+            # Per-class breakdown
+            for cls, cls_m in m["per_class"].items():
+                print(f"    {cls:25s} P:{cls_m['precision']:.3f} "
+                      f"R:{cls_m['recall']:.3f} F1:{cls_m['f1']:.3f}")
+
+        # Compute improvement vs v1 baseline if available
+        print("\n  [Target] Visual-only should improve from ~0.23 acc / 0.08 F1 (v1)")
+        return results
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Multimodal PC Fault Detection Training v2")
+    parser.add_argument("--mode", default="multimodal",
+                        choices=["multimodal", "visual_only", "audio_only"])
+    parser.add_argument("--finetune", default="lora",
+                        choices=["lora", "full", "linear_probe"])
+    parser.add_argument("--epochs", type=int)
+    parser.add_argument("--batch_size", type=int)
+    parser.add_argument("--lr", type=float)
+    parser.add_argument("--fusion", default="concat")
+    parser.add_argument("--modality_dropout", type=float)
+    parser.add_argument("--output_dir", type=str)
+    parser.add_argument("--hub_model_id", type=str)
+    parser.add_argument("--no_push", action="store_true")
+    parser.add_argument("--eval_robustness", action="store_true")
+    parser.add_argument("--quick_test", action="store_true")
+
+    # v2-specific arguments
+    parser.add_argument("--no_ogm", action="store_true",
+                        help="Disable OGM-GE gradient modulation")
+    parser.add_argument("--ogm_alpha", type=float, default=None,
+                        help="OGM-GE modulation strength (default from config)")
+    parser.add_argument("--ogm_noise_sigma", type=float, default=None,
+                        help="OGM-GE noise sigma (default from config)")
+    parser.add_argument("--lambda_visual", type=float, default=None,
+                        help="Visual auxiliary loss weight (default from config)")
+    parser.add_argument("--lambda_audio", type=float, default=None,
+                        help="Audio auxiliary loss weight (default from config)")
+    parser.add_argument("--visual_lr_mult", type=float, default=None,
+                        help="LR multiplier for visual branch (default from config)")
+    parser.add_argument("--audio_lr_mult", type=float, default=None,
+                        help="LR multiplier for audio branch (default from config)")
+
+    args = parser.parse_args()
+
+    # Load config
+    config = ExperimentConfig()
+    config.experiment_name = "multimodal_pc_fault_v2"
+    config.train.mode = args.mode
+    config.train.finetune_method = args.finetune
+    config.model.fusion_type = args.fusion
+
+    if args.epochs:
+        config.train.num_epochs = args.epochs
+    if args.batch_size:
+        config.train.per_device_train_batch_size = args.batch_size
+    if args.lr:
+        config.train.learning_rate = args.lr
+        config.train.lora_learning_rate = args.lr
+    if args.modality_dropout is not None:
+        config.model.modality_dropout_p = args.modality_dropout
+    if args.output_dir:
+        config.train.output_dir = args.output_dir
+    if args.hub_model_id:
+        config.train.hub_model_id = args.hub_model_id
+    if args.no_push:
+        config.train.push_to_hub = False
+    if args.quick_test:
+        config.train.num_epochs = 2
+        config.train.per_device_train_batch_size = 4
+        config.train.per_device_eval_batch_size = 4
+        config.train.gradient_accumulation_steps = 1
+        config.train.logging_steps = 2
+    if args.finetune != "lora":
+        config.lora.enabled = False
+
+    # v2 hyperparameters from config (with CLI overrides)
+    ogm_alpha = args.ogm_alpha if args.ogm_alpha is not None else config.ogm_alpha
+    ogm_noise_sigma = (args.ogm_noise_sigma if args.ogm_noise_sigma is not None
+                       else config.ogm_noise_sigma)
+    lambda_visual = (args.lambda_visual if args.lambda_visual is not None
+                     else config.lambda_visual)
+    lambda_audio = (args.lambda_audio if args.lambda_audio is not None
+                    else config.lambda_audio)
+    visual_lr_mult = (args.visual_lr_mult if args.visual_lr_mult is not None
+                      else config.visual_lr_multiplier)
+    audio_lr_mult = (args.audio_lr_mult if args.audio_lr_mult is not None
+                     else config.audio_lr_multiplier)
+    use_ogm = not args.no_ogm
+
+    # Device and seeds
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    torch.manual_seed(config.train.seed)
+    np.random.seed(config.train.seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(config.train.seed)
+
+    print(f"\n{'='*60}")
+    print(f"Multimodal PC Fault Detection v2")
+    print(f"{'='*60}")
+    print(f"Mode: {args.mode} | Finetune: {args.finetune} | Device: {device}")
+    print(f"OGM-GE: {'ON' if use_ogm else 'OFF'} (alpha={ogm_alpha}, sigma={ogm_noise_sigma})")
+    print(f"Aux loss weights: λ_visual={lambda_visual}, λ_audio={lambda_audio}")
+    print(f"LR multipliers: visual={visual_lr_mult}x, audio={audio_lr_mult}x")
+    print(f"{'='*60}\n")
+
+    # Load processors and dataset
+    vit_proc, ast_ext = get_processors(config.model)
+    train_ds = PCFaultDataset(
+        config.data, config.model, "train", vit_proc, ast_ext, True)
+    val_ds = PCFaultDataset(
+        config.data, config.model, "val", vit_proc, ast_ext, False)
+
+    # Create model
+    model = create_model(
+        config.model, config.lora,
+        mode=args.mode,
+        finetune_method=args.finetune,
+        use_ogm=use_ogm,
+        lambda_visual=lambda_visual,
+        lambda_audio=lambda_audio)
+
+    # Create trainer
+    trainer = MultimodalTrainerV2(
+        model, train_ds, val_ds, config.train, device,
+        use_ogm=use_ogm,
+        ogm_alpha=ogm_alpha,
+        ogm_noise_sigma=ogm_noise_sigma,
+        visual_lr_multiplier=visual_lr_mult,
+        audio_lr_multiplier=audio_lr_mult)
+
+    # Train
+    history = trainer.train()
+
+    # Final evaluation
+    final = trainer.evaluate()
+    print(f"\nFinal Evaluation:")
+    print(f"  Acc={final['accuracy']:.4f} F1={final['macro_f1']:.4f}")
+    for cls, m in final["per_class"].items():
+        print(f"  {cls:25s} P:{m['precision']:.3f} R:{m['recall']:.3f} "
+              f"F1:{m['f1']:.3f} N:{m['support']}")
+
+    # Robustness evaluation
+    robustness_results = None
+    if args.eval_robustness and config.train.mode == "multimodal":
+        robustness_results = trainer.run_robustness_evaluation()
+
+    # Save results
+    os.makedirs(config.train.output_dir, exist_ok=True)
+    results = {
+        "experiment": config.experiment_name,
+        "version": "v2",
+        "mode": config.train.mode,
+        "finetune_method": config.train.finetune_method,
+        "anti_collapse_config": {
+            "ogm_ge": use_ogm,
+            "ogm_alpha": ogm_alpha,
+            "ogm_noise_sigma": ogm_noise_sigma,
+            "lambda_visual": lambda_visual,
+            "lambda_audio": lambda_audio,
+            "visual_lr_multiplier": visual_lr_mult,
+            "audio_lr_multiplier": audio_lr_mult,
+        },
+        "final_metrics": {
+            "accuracy": final["accuracy"],
+            "macro_f1": final["macro_f1"],
+            "weighted_f1": final["weighted_f1"],
+            "per_class": final["per_class"],
+            "confusion_matrix": final["confusion_matrix"],
+        },
+        "history": history,
+        "best_epoch": trainer.best_epoch,
+        "best_metric": trainer.best_metric,
+    }
+    if robustness_results:
+        results["robustness"] = robustness_results
+
+    with open(os.path.join(config.train.output_dir, "results_v2.json"), "w") as f:
+        json.dump(results, f, indent=2)
+    print(f"\nResults saved to {config.train.output_dir}/results_v2.json")
+
+    # Push to hub
+    if config.train.push_to_hub:
+        try:
+            from huggingface_hub import HfApi, login
+            login(token=os.environ.get("HF_TOKEN"))
+            HfApi().upload_folder(
+                folder_path=config.train.output_dir,
+                repo_id=config.train.hub_model_id,
+                repo_type="model",
+                commit_message=f"Training v2: {config.experiment_name} (OGM-GE)")
+            print(f"✓ Pushed to https://huggingface.co/{config.train.hub_model_id}")
+        except Exception as e:
+            print(f"✗ Push failed: {e}")
+
+
+if __name__ == "__main__":
+    main()