v4 update

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