Add training script and ablation runner

src/train.py

@@ -0,0 +1,214 @@
+"""
+Multimodal PC Fault Detection - Training Script
+=================================================
+Usage:
+  python train.py --mode multimodal --finetune lora --eval_robustness
+  python train.py --mode visual_only --finetune lora --no_push
+  python train.py --mode audio_only --finetune lora --no_push
+  python train.py --mode multimodal --finetune full --lr 2e-5
+  python train.py --quick_test --no_push
+"""
+
+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 import PCFaultDataset, multimodal_collate_fn
+from models import create_model, get_processors
+
+
+def compute_metrics(preds, labels, class_names=FAULT_CLASSES):
+    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 MultimodalTrainer:
+    def __init__(self, model, train_dataset, val_dataset, config, device):
+        self.model = model.to(device)
+        self.device, self.config = device, config
+        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)
+        param_groups = self._get_param_groups()
+        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=total_steps, pct_start=config.warmup_ratio, anneal_strategy="cos")
+        self.scaler = torch.amp.GradScaler("cuda") if config.fp16 and device.type == "cuda" else None
+        self.best_metric, self.best_epoch = 0.0, 0
+        self.history = {"train_loss": [], "val_loss": [], "val_accuracy": [], "val_macro_f1": []}
+    
+    def _get_param_groups(self):
+        lora_params, other_params = [], []
+        for name, param in self.model.named_parameters():
+            if not param.requires_grad: continue
+            (lora_params if "lora" in name.lower() else other_params).append(param)
+        groups = []
+        if lora_params: groups.append({"params": lora_params, "lr": self.config.lora_learning_rate})
+        if other_params: groups.append({"params": other_params, "lr": self.config.learning_rate})
+        if not groups: raise ValueError("No trainable parameters!")
+        return groups
+    
+    def train_epoch(self, epoch):
+        self.model.train()
+        total_loss, num_batches = 0.0, 0
+        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 (batch_idx + 1) % self.config.gradient_accumulation_steps == 0:
+                if self.scaler:
+                    self.scaler.unscale_(self.optimizer)
+                    torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.config.max_grad_norm)
+                    self.scaler.step(self.optimizer)
+                    self.scaler.update()
+                else:
+                    torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.config.max_grad_norm)
+                    self.optimizer.step()
+                self.scheduler.step()
+                self.optimizer.zero_grad()
+            if (batch_idx + 1) % self.config.logging_steps == 0 or batch_idx == 0:
+                print(f"  [Epoch {epoch+1}] Step {batch_idx+1}/{len(self.train_loader)} | Loss: {total_loss/num_batches:.4f} | LR: {self.optimizer.param_groups[0]['lr']:.2e}")
+        return total_loss / max(num_batches, 1)
+    
+    @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"\\nTraining: mode={self.model.mode}, epochs={self.config.num_epochs}, batch={self.config.per_device_train_batch_size}, device={self.device}")
+        for epoch in range(self.config.num_epochs):
+            t0 = time.time()
+            train_loss = self.train_epoch(epoch)
+            val_metrics = self.evaluate()
+            print(f"\\n[Epoch {epoch+1}/{self.config.num_epochs}] ({time.time()-t0:.1f}s) Train Loss: {train_loss:.4f} | Val Loss: {val_metrics['val_loss']:.4f} | Acc: {val_metrics['accuracy']:.4f} | F1: {val_metrics['macro_f1']:.4f}")
+            self.history["train_loss"].append(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"])
+            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 ===")
+        results = {}
+        for name, mask in [("both", 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}: Acc={m['accuracy']:.4f} F1={m['macro_f1']:.4f}")
+        return results
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    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")
+    args = parser.parse_args()
+    
+    config = ExperimentConfig()
+    config.train.mode, config.train.finetune_method, config.model.fusion_type = args.mode, args.finetune, 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
+    
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    torch.manual_seed(config.train.seed); np.random.seed(config.train.seed)
+    
+    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)
+    model = create_model(config.model, config.lora, config.train.mode, config.train.finetune_method)
+    trainer = MultimodalTrainer(model, train_ds, val_ds, config.train, device)
+    history = trainer.train()
+    
+    final = trainer.evaluate()
+    print(f"\\nFinal: 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']}")
+    
+    if args.eval_robustness and config.train.mode == "multimodal":
+        trainer.run_robustness_evaluation()
+    
+    os.makedirs(config.train.output_dir, exist_ok=True)
+    with open(os.path.join(config.train.output_dir, "results.json"), "w") as f:
+        json.dump({"experiment": config.experiment_name, "mode": config.train.mode, "finetune_method": config.train.finetune_method,
+            "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}, f, indent=2)
+    
+    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: {config.experiment_name}")
+            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()