v4 update

v4/classes/towers/clinical_tower.py

@@ -0,0 +1,94 @@
+"""clinical_tower — ClinicalEncoder for v4.
+
+Self-contained: no v3 dependencies.
+Inherits get_sample dispatch from TowerBase.
+"""
+from __future__ import annotations
+
+import numpy as np
+import torch
+from torch import nn
+
+from v4.classes.towerbase import TowerBase
+from v4.classes.accessory.se_block import SEBlock
+
+
+class ClinicalEncoder(TowerBase):
+    """MLP over tabular clinical features.
+
+    clinical_data : ClinicalDataView — provides feature_dim, vectorize_entity, side_map
+    hidden_dim    : output embedding dimensionality
+    dropout       : applied after the first linear block
+    use_se        : wrap output with SEBlock channel gating
+    se_reduction  : SEBlock bottleneck factor
+    se_pre_norm   : apply LayerNorm before SEBlock
+    """
+
+    def __init__(
+        self,
+        clinical_data,
+        hidden_dim: int = 128,
+        dropout: float = 0.1,
+        use_se: bool = False,
+        se_reduction: int = 16,
+        se_pre_norm: bool = True,
+    ):
+        super().__init__()
+        self.clinical_data = clinical_data
+        self._out_dim      = hidden_dim
+        feature_dim        = clinical_data.feature_dim
+
+        self.block0 = nn.Sequential(
+            nn.Linear(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
+
+    # ── TowerBase interface ──────────────────────────────────────────────────
+
+    @property
+    def out_dim(self) -> int:
+        return self._out_dim
+
+    @property
+    def _side_map(self) -> dict[str, str]:
+        return self.clinical_data.side_map
+
+    def _get(self, *ids) -> torch.Tensor:
+        arr = self.clinical_data.vectorize_entity(*ids)
+        return torch.from_numpy(arr.astype(np.float32, copy=False))
+
+    # ── nn.Module forward ────────────────────────────────────────────────────
+
+    def forward(self, x) -> torch.Tensor:
+        if not isinstance(x, torch.Tensor):
+            x = torch.as_tensor(x, dtype=torch.float32)
+        h = self.net(x)
+        if self.tower_se is not None:
+            h, _ = self.tower_se(self.tower_ln(h))
+        return h
+
+    # ── Utilities ────────────────────────────────────────────────────────────
+
+    def set_freeze_ratio(self, ratio: float) -> None:
+        """Freeze the earliest MLP block proportionally."""
+        r = max(0.0, min(1.0, float(ratio)))
+        for p in self.block0.parameters():
+            p.requires_grad = True
+        for p in self.block1.parameters():
+            p.requires_grad = True
+        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

v4/classes/towers/image_tower.py

@@ -0,0 +1,83 @@
+"""image_tower — ImageEncoder for v4.
+
+Self-contained: no v3 dependencies.
+Inherits get_sample dispatch from TowerBase.
+"""
+from __future__ import annotations
+
+import math
+
+import torch
+from torch import nn
+
+from v4.classes.towerbase import TowerBase
+from v4.classes.accessory.backbones import build_backbone
+from v4.classes.accessory.se_block import SEBlock
+from v4.classes.accessory.transforms import build_backbone_transform, build_eval_transform
+
+
+class ImageEncoder(TowerBase):
+    """Vision backbone → pooled feature vector.
+
+    image_data   : ImageDataView — provides load_image(*ids) and side_map
+    backbone     : backbone key (see accessory/backbones.py)
+    freeze_ratio : fraction of early blocks to freeze in [0, 1]
+    use_se       : apply SE attention over the pooled feature vector
+    augment      : include random flip/rotation/jitter in the train transform
+    """
+
+    def __init__(
+        self,
+        image_data,
+        backbone:     str   = "efficientnet_b0",
+        freeze_ratio: float = 0.0,
+        use_se:       bool  = False,
+        se_reduction: int   = 16,
+        se_pre_norm:  bool  = True,
+        augment:      bool  = True,
+    ):
+        super().__init__()
+        self.image_data     = image_data
+        self._name          = backbone
+        self.backbone, self._base_dim, self._blocks = build_backbone(backbone, freeze_ratio)
+        self.transform      = build_backbone_transform(backbone, augment=augment)
+        self.eval_transform = build_eval_transform(backbone)
+
+        self.tower_ln = nn.LayerNorm(self._base_dim) if se_pre_norm else nn.Identity()
+        self.tower_se = SEBlock(self._base_dim, reduction=se_reduction, residual=True) if use_se else None
+
+    # ── TowerBase interface ──────────────────────────────────────────────────
+
+    @property
+    def out_dim(self) -> int:
+        return self._base_dim
+
+    @property
+    def _side_map(self) -> dict[str, str]:
+        return self.image_data.side_map
+
+    def _get(self, *ids) -> torch.Tensor:
+        img = self.image_data.load_image(*ids)
+        t   = self.transform if self.training else self.eval_transform
+        return t(img)
+
+    # ── nn.Module forward ────────────────────────────────────────────────────
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        y = self.backbone(x)
+        if self.tower_se is not None:
+            y, _ = self.tower_se(self.tower_ln(y))
+        return y
+
+    # ── Utilities ────────────────────────────────────────────────────────────
+
+    def set_freeze_ratio(self, ratio: float) -> None:
+        """Dynamically freeze the earliest floor(N * ratio) backbone blocks."""
+        r        = max(0.0, min(1.0, float(ratio)))
+        n_freeze = int(math.floor(len(self._blocks) * r))
+        for b in self._blocks:
+            for p in b.parameters():
+                p.requires_grad = True
+        for b in self._blocks[:n_freeze]:
+            for p in b.parameters():
+                p.requires_grad = False