""" 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) - Supports both old proxy data (dataset_real) and new built data (dataset_v2) Usage: # With old proxy data (ToyADMOS + MVTec, default) python train_v2.py --mode multimodal --finetune lora --eval_robustness # With new built dataset (from build_dataset.py) python train_v2.py --dataset local --dataset_dir ../data/dataset_build --eval_robustness python train_v2.py --dataset hub --hub_dataset Ellaft/pc-fault-real-dataset # Other options python train_v2.py --mode visual_only --finetune lora --no_push 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 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, collate_fn=None): 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=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=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": [], "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): visual_params, audio_params, fusion_params = [], [], [] 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) base_lr = self.config.lora_learning_rate 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): self.model.train() total_loss, total_loss_fusion, total_loss_visual, total_loss_audio = 0.0, 0.0, 0.0, 0.0 num_batches = 0 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) 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() 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"] 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"]) if (batch_idx + 1) % self.config.gradient_accumulation_steps == 0: if self.scaler: self.scaler.unscale_(self.optimizer) if (self.use_ogm and self.ogm is not None and ogm_cv): 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() 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} | 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} c_a: {ogm_ca[-1]:.3f}" print(msg) 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.update({"ogm_visual_conf": np.mean(ogm_v_confs), "ogm_audio_conf": np.mean(ogm_a_confs), "ogm_coeff_visual": np.mean(ogm_cv), "ogm_coeff_audio": np.mean(ogm_ca)}) return epoch_stats @torch.no_grad() def evaluate(self, modality_mask=None): self.model.eval() all_preds, all_labels, total_loss, num_batches = [], [], 0.0, 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): 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}, λ_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() 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} audio={train_stats['loss_audio']:.4f})") print(loss_msg) print(f" Val Loss: {val_metrics['val_loss']:.4f} | Acc: {val_metrics['accuracy']:.4f} | 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} audio_conf={train_stats['ogm_audio_conf']:.4f} " f"| coeff_v={train_stats['ogm_coeff_visual']:.4f} coeff_a={train_stats['ogm_coeff_audio']:.4f}") 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"]) 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}, Best F1={self.best_metric:.4f}") return self.history def run_robustness_evaluation(self): print("\n=== Missing Modality Robustness Evaluation ===") results = {} for name, mask in [("both_modalities", None), ("visual_only", {"visual": 1.0, "audio": 0.0}), ("audio_only", {"visual": 0.0, "audio": 1.0})]: 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}") for cls, cls_m in m["per_class"].items(): print(f" {cls:25s} P:{cls_m['precision']:.3f} R:{cls_m['recall']:.3f} F1:{cls_m['f1']:.3f}") 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") # Dataset selection parser.add_argument("--dataset", default="proxy", choices=["proxy", "local", "hub"], help="Dataset source: 'proxy' (ToyADMOS+MVTec, default), " "'local' (build_dataset.py output), 'hub' (HF Hub dataset)") parser.add_argument("--dataset_dir", default="./dataset_build", help="Path to build_dataset.py output (for --dataset local)") parser.add_argument("--hub_dataset", default="Ellaft/pc-fault-real-dataset", help="HuggingFace dataset ID (for --dataset hub)") # v2-specific arguments parser.add_argument("--no_ogm", action="store_true") parser.add_argument("--ogm_alpha", type=float, default=None) parser.add_argument("--ogm_noise_sigma", type=float, default=None) parser.add_argument("--lambda_visual", type=float, default=None) parser.add_argument("--lambda_audio", type=float, default=None) parser.add_argument("--visual_lr_mult", type=float, default=None) parser.add_argument("--audio_lr_mult", type=float, default=None) 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 = 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, config.train.per_device_train_batch_size = 2, 4 config.train.per_device_eval_batch_size, config.train.gradient_accumulation_steps = 4, 1 config.train.logging_steps = 2 if args.finetune != "lora": config.lora.enabled = False 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 = 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"Dataset: {args.dataset}" + (f" ({args.dataset_dir})" if args.dataset == "local" else f" ({args.hub_dataset})" if args.dataset == "hub" else " (ToyADMOS + MVTec proxy)")) 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 vit_proc, ast_ext = get_processors(config.model) # ---- Load dataset based on --dataset flag ---- if args.dataset in ("local", "hub"): from dataset_v2 import BuiltDataset as PCFaultDataset, multimodal_collate_fn source = args.dataset # "local" or "hub" train_ds = PCFaultDataset( config.data, config.model, "train", vit_proc, ast_ext, True, source=source, dataset_dir=args.dataset_dir, hub_dataset=args.hub_dataset) val_ds = PCFaultDataset( config.data, config.model, "val", vit_proc, ast_ext, False, source=source, dataset_dir=args.dataset_dir, hub_dataset=args.hub_dataset) else: # Default: old proxy data (ToyADMOS + MVTec) from dataset_real import RealPCFaultDataset as PCFaultDataset, multimodal_collate_fn 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, collate_fn=multimodal_collate_fn) # 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} F1:{m['f1']:.3f} N:{m['support']}") 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, "dataset_source": args.dataset, "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") 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()