Add distributed server implementation and protocol definitions

- Introduced `protocol.py` for shared data models used in server/client communication, including request and response schemas for registration, job submission, and status updates. - Implemented `server.py` to manage a SQLite job queue and client registry, handling job polling, status updates, and job completion. - Created a cheat sheet for server usage, detailing commands for starting the server, submitting jobs, and monitoring clients. - Added several experiment configuration files for various training setups, including geometry vector injections and baseline ensembles.
2026-04-28 08:24:25 +02:00
parent 4dea45df78
commit 512ebd13b2
42 changed files with 4468 additions and 612 deletions
@@ -2,50 +2,128 @@

 Self-contained: no v3 dependencies.
 Inherits get_sample dispatch from TowerBase.
+
+Geometry injection (EPC supply)
+--------------------------------
+When geometry_source is set, ImageEncoder asks the image_data view for a loader
+via image_data.build_geometry_loader(source, **kwargs).  The view is responsible
+for understanding what that source means for its specific domain (fundus contours,
+U-Net segmentations, cat ear landmarks, etc.).
+
+During early_pass the loader pre-computes all per-entity geometry vectors and
+publishes them to the EarlyPassContext under the key "geometry_vectors"
+({(entity_id...): np.ndarray of length geom_dim}).  ClinicalEncoder (or any
+other tower with epc_requests: ["geometry_vectors"]) can then consume them.
+
+The tower reads feature_dim and feature_names from the loader instance, so it
+can log geometry info without knowing anything about CDR, disc masks, or other
+domain-specific concepts.
+
+Config example:
+{
+  "name":         "img",
+  "module":       "v4.classes.towers.image_tower",
+  "class":        "ImageEncoder",
+  "data_source":  "image",
+  "epc_supplies": ["geometry_vectors"],
+  "args": {
+    "backbone":        "refugelike",
+    "augment":         true,
+    "geometry_source": "gt",
+    "contour_dir":     "Papila/ExpertsSegmentations/Contours"
+  }
+}
 """
 from __future__ import annotations

 import math
+import sys
+from pathlib import Path
+from typing import Any

 import torch
 from torch import nn

+_REPO_ROOT = Path(__file__).resolve().parents[3]
+if str(_REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(_REPO_ROOT))
+
 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
+from v4.classes.accessory.transforms import (
+    build_backbone_transform, build_eval_transform, build_split_transforms,
+)


 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
+    image_data        : ImageDataView — provides load_image(*ids) and side_map.
+                        Must implement build_geometry_loader(source, **kwargs)
+                        if geometry_source is set.
+    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
+    cache_transformed : if True, cache resized + ToTensor'd float32 [0, 1] CHW
+                        tensors per fold.  Per-batch cost drops to augment +
+                        Normalize on tensors only (no PIL, no Resize, no decode).
+                        Memory: ~3 × crop_size² × 4B per cached image.
+                        Cache is rebuilt at the start of every fold via early_pass.
+    geometry_source   : source key passed to image_data.build_geometry_loader()
+                        (e.g. "gt", "unet").  None = geometry disabled.
+    **geom_kwargs     : forwarded verbatim to build_geometry_loader() — e.g.
+                        contour_dir="Papila/ExpertsSegmentations/Contours"
    """

+    EPC_GEOMETRY_KEY = "geometry_vectors"
+
    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,
+        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,
+        cache_transformed: bool       = False,
+        geometry_source:   str | None = None,
+        **geom_kwargs: Any,
    ):
        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._cache_transformed = cache_transformed
+        if cache_transformed:
+            self._precache_tf, self._post_train_tf = build_split_transforms(backbone, augment=augment)
+            _,                 self._post_eval_tf  = build_split_transforms(backbone, augment=False)
+            self._tensor_cache: dict[tuple, torch.Tensor] = {}
+        else:
+            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

+        self._geom_loader = None
+        if geometry_source is not None:
+            if not hasattr(image_data, "build_geometry_loader"):
+                raise TypeError(
+                    f"ImageEncoder geometry_source={geometry_source!r} requires "
+                    f"image_data to implement build_geometry_loader(), "
+                    f"but {type(image_data).__name__} does not."
+                )
+            self._geom_loader = image_data.build_geometry_loader(geometry_source, **geom_kwargs)
+            print(
+                f"[ImageEncoder] geometry_source={geometry_source!r}  "
+                f"features={self._geom_loader.feature_names}",
+                flush=True,
+            )
+
    # ── TowerBase interface ──────────────────────────────────────────────────

    @property
@@ -57,10 +135,61 @@ class ImageEncoder(TowerBase):
        return self.image_data.side_map

    def _get(self, *ids) -> torch.Tensor:
+        if self._cache_transformed:
+            key = tuple(ids)
+            cached = self._tensor_cache.get(key)
+            if cached is None:
+                cached = self._precache_tf(self.image_data.load_image(*ids))
+                self._tensor_cache[key] = cached
+            tail = self._post_train_tf if self.training else self._post_eval_tf
+            return tail(cached)
        img = self.image_data.load_image(*ids)
        t   = self.transform if self.training else self.eval_transform
        return t(img)

+    # ── EPC early_pass ───────────────────────────────────────────────────────
+
+    def early_pass(self, context) -> None:
+        """Per-fold setup: warm tensor cache (if enabled), publish geometry vectors."""
+        data = context.require("data")
+
+        if self._cache_transformed:
+            self._tensor_cache.clear()
+            n = self._warm_tensor_cache(data, context.require("split"))
+            print(
+                f"[ImageEncoder] warmed transformed-tensor cache for {n} entries "
+                f"({self._name})",
+                flush=True,
+            )
+
+        if self._geom_loader is None:
+            return
+        self._geom_loader.precompute(data.df, patient_col=data.patient_col)
+        vecs = self._geom_loader.all_vectors()
+        context.put(self.EPC_GEOMETRY_KEY, vecs)
+        print(
+            f"[ImageEncoder] published {len(vecs)} geometry vectors "
+            f"(dim={self._geom_loader.feature_dim}) to EPC key '{self.EPC_GEOMETRY_KEY}'",
+            flush=True,
+        )
+
+    def _warm_tensor_cache(self, data, split) -> int:
+        """Pre-fill the per-tower tensor cache for all entries in this fold's splits."""
+        seen: set[tuple] = set()
+        for df in (split.train, split.val, split.test):
+            if df is None or len(df) == 0:
+                continue
+            pc = data.patient_col
+            for _, row in df.iterrows():
+                pid = int(row[pc])
+                eye = str(row.get("eyeID", "OD"))
+                key = (pid, eye)
+                if key in self._tensor_cache or key in seen:
+                    continue
+                self._tensor_cache[key] = self._precache_tf(self.image_data.load_image(pid, eye))
+                seen.add(key)
+        return len(self._tensor_cache)
+
    # ── nn.Module forward ────────────────────────────────────────────────────

    def forward(self, x: torch.Tensor) -> torch.Tensor: