moved_repo_first_update

classes/md_tower.py

@@ -0,0 +1,54 @@
+# md_tower.py
+import torch
+import torch.nn as nn
+from classes import ClinicalData
+from classes.SE_attention import SEBlock
+
+class MDTower(nn.Module):
+    """MLP over ClinicalData.vectorize_row outputs (convert to torch inside tower)."""
+    def __init__(self, clinical_data: ClinicalData, hidden_dim: int = 128, dropout: float = 0.1,
+                 use_se: bool = False, se_reduction: int = 16, se_pre_norm: bool = True):
+        super().__init__()
+        self.feature_dim = clinical_data.feature_dim
+        self.out_dim = hidden_dim
+        # two-block MLP so we can optionally freeze/thaw per block
+        self.block0 = nn.Sequential(
+            nn.Linear(self.feature_dim, hidden_dim),
+            nn.LayerNorm(hidden_dim),
+            nn.ReLU(inplace=True),
+            nn.Dropout(dropout),
+        )
+        self.block1 = nn.Sequential(
+            nn.Linear(hidden_dim, hidden_dim),
+            nn.ReLU(inplace=True),
+        )
+        self.net = nn.Sequential(self.block0, self.block1)
+        self.tower_ln = nn.LayerNorm(hidden_dim) if se_pre_norm else nn.Identity()
+        self.tower_se = SEBlock(hidden_dim, reduction=se_reduction, residual=True) if use_se else None
+
+    def forward(self, meta_np_or_torch) -> torch.Tensor:
+        if isinstance(meta_np_or_torch, torch.Tensor):
+            x = meta_np_or_torch
+        else:
+            x = torch.as_tensor(meta_np_or_torch, dtype=torch.float32)
+        h = self.net(x)
+        if self.tower_se is not None:
+            h, _ = self.tower_se(self.tower_ln(h))
+        return h
+
+    def set_freeze_ratio(self, ratio: float):
+        """Optionally freeze earliest blocks of the MLP.
+        With two blocks, ratio≥0.5 freezes block0; ratio≥1.0 freezes both."""
+        r = max(0.0, min(1.0, float(ratio)))
+        # Unfreeze all
+        for p in self.block0.parameters():
+            p.requires_grad = True
+        for p in self.block1.parameters():
+            p.requires_grad = True
+        # Freeze earliest blocks based on ratio threshold
+        if r >= 0.5:
+            for p in self.block0.parameters():
+                p.requires_grad = False
+        if r >= 1.0:
+            for p in self.block1.parameters():
+                p.requires_grad = False