Add model architecture (ViT + AST + Late Fusion + LoRA)

src/models.py

@@ -0,0 +1,131 @@
+"""
+Multimodal PC Fault Detection - Model Architecture
+====================================================
+Two-branch architecture:
+  - Visual: ViT-B/16 pretrained on ImageNet-21k
+  - Audio: AST pretrained on AudioSet
+  - Fusion: Late fusion (concat / weighted sum / attention)
+  
+Supports LoRA, full fine-tuning, and linear probe modes.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Dict, Optional, Literal
+from transformers import ViTModel, ASTModel, ViTImageProcessor, ASTFeatureExtractor
+from peft import LoraConfig, get_peft_model
+from config import ModelConfig, LoRAConfig, FAULT_CLASSES
+
+
+class VisualBranch(nn.Module):
+    def __init__(self, config, lora_config=None, finetune_method="lora"):
+        super().__init__()
+        self.vit = ViTModel.from_pretrained(config.vit_model_name)
+        if finetune_method == "lora" and lora_config and lora_config.enabled:
+            peft_config = LoraConfig(r=lora_config.r, lora_alpha=lora_config.lora_alpha,
+                target_modules=lora_config.vit_target_modules, lora_dropout=lora_config.lora_dropout, bias=lora_config.bias)
+            self.vit = get_peft_model(self.vit, peft_config)
+            self.vit.print_trainable_parameters()
+        elif finetune_method == "linear_probe":
+            for param in self.vit.parameters(): param.requires_grad = False
+    
+    def forward(self, pixel_values):
+        return self.vit(pixel_values=pixel_values).last_hidden_state[:, 0, :]
+
+
+class AudioBranch(nn.Module):
+    def __init__(self, config, lora_config=None, finetune_method="lora"):
+        super().__init__()
+        self.ast = ASTModel.from_pretrained(config.ast_model_name)
+        if finetune_method == "lora" and lora_config and lora_config.enabled:
+            peft_config = LoraConfig(r=lora_config.r, lora_alpha=lora_config.lora_alpha,
+                target_modules=lora_config.ast_target_modules, lora_dropout=lora_config.lora_dropout, bias=lora_config.bias)
+            self.ast = get_peft_model(self.ast, peft_config)
+            self.ast.print_trainable_parameters()
+        elif finetune_method == "linear_probe":
+            for param in self.ast.parameters(): param.requires_grad = False
+    
+    def forward(self, input_values):
+        return self.ast(input_values=input_values).last_hidden_state[:, 0, :]
+
+
+class LateFusion(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.fusion_type = config.fusion_type
+        if config.fusion_type == "concat":
+            self.visual_proj = nn.Linear(config.vit_embed_dim, config.fusion_dim)
+            self.audio_proj = nn.Linear(config.ast_embed_dim, config.fusion_dim)
+            self.classifier = nn.Sequential(nn.LayerNorm(config.fusion_dim * 2), nn.Dropout(config.fusion_dropout),
+                nn.Linear(config.fusion_dim * 2, config.fusion_dim), nn.GELU(), nn.Dropout(config.fusion_dropout),
+                nn.Linear(config.fusion_dim, config.num_classes))
+        elif config.fusion_type == "weighted_sum":
+            self.visual_head = nn.Linear(config.vit_embed_dim, config.num_classes)
+            self.audio_head = nn.Linear(config.ast_embed_dim, config.num_classes)
+            self.fusion_weights = nn.Parameter(torch.tensor([0.5, 0.5]))
+        elif config.fusion_type == "attention":
+            self.visual_proj = nn.Linear(config.vit_embed_dim, config.fusion_dim)
+            self.audio_proj = nn.Linear(config.ast_embed_dim, config.fusion_dim)
+            self.cross_attn = nn.MultiheadAttention(embed_dim=config.fusion_dim, num_heads=8, dropout=config.fusion_dropout, batch_first=True)
+            self.classifier = nn.Sequential(nn.LayerNorm(config.fusion_dim), nn.Dropout(config.fusion_dropout), nn.Linear(config.fusion_dim, config.num_classes))
+    
+    def forward(self, visual_emb, audio_emb, modality_mask=None):
+        if modality_mask:
+            visual_emb = visual_emb * modality_mask.get("visual", 1.0)
+            audio_emb = audio_emb * modality_mask.get("audio", 1.0)
+        if self.fusion_type == "concat":
+            fused = torch.cat([self.visual_proj(visual_emb), self.audio_proj(audio_emb)], dim=-1)
+            return self.classifier(fused)
+        elif self.fusion_type == "weighted_sum":
+            w = torch.softmax(self.fusion_weights, dim=0)
+            return w[0] * self.visual_head(visual_emb) + w[1] * self.audio_head(audio_emb)
+        elif self.fusion_type == "attention":
+            tokens = torch.cat([self.visual_proj(visual_emb).unsqueeze(1), self.audio_proj(audio_emb).unsqueeze(1)], dim=1)
+            return self.classifier(self.cross_attn(tokens, tokens, tokens)[0].mean(dim=1))
+
+
+class MultimodalPCFaultDetector(nn.Module):
+    def __init__(self, model_config, lora_config=None, finetune_method="lora", mode="multimodal"):
+        super().__init__()
+        self.mode, self.modality_dropout_p = mode, model_config.modality_dropout_p
+        self.visual_branch = VisualBranch(model_config, lora_config, finetune_method) if mode in ("multimodal", "visual_only") else None
+        self.audio_branch = AudioBranch(model_config, lora_config, finetune_method) if mode in ("multimodal", "audio_only") else None
+        if mode == "multimodal":
+            self.fusion = LateFusion(model_config)
+        else:
+            embed_dim = model_config.vit_embed_dim if mode == "visual_only" else model_config.ast_embed_dim
+            self.classifier = nn.Sequential(nn.LayerNorm(embed_dim), nn.Dropout(model_config.fusion_dropout),
+                nn.Linear(embed_dim, model_config.fusion_dim), nn.GELU(), nn.Dropout(model_config.fusion_dropout),
+                nn.Linear(model_config.fusion_dim, model_config.num_classes))
+        self.loss_fn = nn.CrossEntropyLoss()
+        total = sum(p.numel() for p in self.parameters())
+        trainable = sum(p.numel() for p in self.parameters() if p.requires_grad)
+        print(f"[Model] Mode={mode}, Total={total:,}, Trainable={trainable:,} ({100*trainable/total:.2f}%)")
+    
+    def forward(self, pixel_values=None, audio_values=None, labels=None):
+        if self.mode == "multimodal":
+            v_emb, a_emb = self.visual_branch(pixel_values), self.audio_branch(audio_values)
+            mask = None
+            if self.training and self.modality_dropout_p > 0:
+                mask = {"visual": 0.0 if torch.rand(1).item() < self.modality_dropout_p else 1.0,
+                        "audio": 0.0 if torch.rand(1).item() < self.modality_dropout_p else 1.0}
+                if mask["visual"] == 0.0 and mask["audio"] == 0.0:
+                    mask["visual" if torch.rand(1).item() < 0.5 else "audio"] = 1.0
+            logits = self.fusion(v_emb, a_emb, mask)
+        elif self.mode == "visual_only":
+            logits = self.classifier(self.visual_branch(pixel_values))
+        else:
+            logits = self.classifier(self.audio_branch(audio_values))
+        outputs = {"logits": logits}
+        if labels is not None:
+            outputs["loss"] = self.loss_fn(logits, labels)
+        return outputs
+
+
+def create_model(model_config, lora_config, mode="multimodal", finetune_method="lora"):
+    return MultimodalPCFaultDetector(model_config, lora_config, finetune_method, mode)
+
+def get_processors(model_config):
+    return (ViTImageProcessor.from_pretrained(model_config.vit_model_name),
+            ASTFeatureExtractor.from_pretrained(model_config.ast_model_name))