Add analysis scripts and experiment configurations for bridge attention and sensitivity studies

- Introduced `bridge_attention_ceiling_check.py` for variance decomposition analysis on bridge attention configurations.
- Added `bridge_attention_readout.py` to perform per-tower gate and contribution readouts, including AUC sanity checks.
- Created multiple JSON configuration files for backbone replication experiments, including anonymous CV variants and basic backbones.
- Implemented sensitivity experiments to evaluate the impact of axial length inclusion and EfficientNetV2-M performance at higher resolutions.
- Added a memory probe script to assess GPU memory usage during training with EfficientNetV2-M.
This commit is contained in:
rpotter6298
2026-07-03 08:51:44 +02:00
parent 3d954a4606
commit 708fbc70ce
52 changed files with 2223 additions and 218 deletions
+53 -17
View File
@@ -76,39 +76,75 @@ class ImageTransformConfig:
return transforms.Compose(ops)
def backbone_transform_config(backbone_name: str, augment: bool = True) -> ImageTransformConfig:
"""Build an ImageTransformConfig using the backbone's default normalisation stats."""
def backbone_transform_config(
backbone_name: str,
augment: bool = True,
crop_size: int | None = None,
resize_size: int | None = None,
) -> ImageTransformConfig:
"""Build an ImageTransformConfig using the backbone's default normalisation stats.
crop_size / resize_size override the backbone's default input resolution. When
crop_size is overridden but resize_size is not, resize_size is scaled
proportionally (8/7 ratio, matching the standard 224 → 256 pattern).
"""
key = (backbone_name or "").lower()
if _is_timm_backbone(key):
# ConvNeXt-V2 and other timm models we currently expose are all
# pretrained with standard ImageNet stats at 224×224.
return ImageTransformConfig(crop_size=224, mean=IMAGENET_MEAN,
std=IMAGENET_STD, augment=augment)
if key not in BACKBONES:
raise ValueError(f"Unknown backbone '{backbone_name}'.")
spec = BACKBONES[key]
mean = getattr(spec.weights_default, "meta", {}).get("mean", IMAGENET_MEAN)
std = getattr(spec.weights_default, "meta", {}).get("std", IMAGENET_STD)
crop = 299 if key == "inception_v3" else 224
return ImageTransformConfig(crop_size=crop, mean=mean, std=std, augment=augment)
mean, std = IMAGENET_MEAN, IMAGENET_STD
default_crop = 224
else:
if key not in BACKBONES:
raise ValueError(f"Unknown backbone '{backbone_name}'.")
spec = BACKBONES[key]
mean = getattr(spec.weights_default, "meta", {}).get("mean", IMAGENET_MEAN)
std = getattr(spec.weights_default, "meta", {}).get("std", IMAGENET_STD)
default_crop = 299 if key == "inception_v3" else 224
crop = crop_size if crop_size is not None else default_crop
resize = resize_size if resize_size is not None else round(crop * 8 / 7)
return ImageTransformConfig(crop_size=crop, resize_size=resize,
mean=mean, std=std, augment=augment)
def build_backbone_transform(backbone_name: str, augment: bool = True) -> transforms.Compose:
return backbone_transform_config(backbone_name, augment=augment).build()
def build_backbone_transform(
backbone_name: str,
augment: bool = True,
crop_size: int | None = None,
resize_size: int | None = None,
) -> transforms.Compose:
return backbone_transform_config(
backbone_name, augment=augment,
crop_size=crop_size, resize_size=resize_size,
).build()
def build_eval_transform(backbone_name: str) -> transforms.Compose:
def build_eval_transform(
backbone_name: str,
crop_size: int | None = None,
resize_size: int | None = None,
) -> transforms.Compose:
"""Deterministic eval transform — no augmentation, backbone-matched normalisation."""
return build_backbone_transform(backbone_name, augment=False)
return build_backbone_transform(
backbone_name, augment=False,
crop_size=crop_size, resize_size=resize_size,
)
def build_split_transforms(
backbone_name: str, augment: bool = True
backbone_name: str,
augment: bool = True,
crop_size: int | None = None,
resize_size: int | None = None,
) -> tuple[transforms.Compose, transforms.Compose]:
"""Return (precache, postcache) transform pair for tensor-cached image towers.
precache : PIL → CHW float32 in [0, 1] (deterministic, run once at fill)
postcache : tensor → augmented + normalized tensor (run per batch)
"""
cfg = backbone_transform_config(backbone_name, augment=augment)
cfg = backbone_transform_config(
backbone_name, augment=augment,
crop_size=crop_size, resize_size=resize_size,
)
return cfg.build_precache(), cfg.build_postcache()
+26 -14
View File
@@ -6,19 +6,19 @@ FeatureStore — records embeddings (opt-in); same structure but per-head
HDF5 layout — PredictionStore
------------------------------
/{phase}/logits float32 (n_folds, n_epochs, n_samples, n_heads, n_classes)
/{phase}/head_names str (n_heads,)
/{phase}/y_true int64 (n_samples,)
/{phase}/entity_id_{k} int64|str (n_samples,) — one dataset per id component
/{phase}/split str (n_folds, n_samples)
/{phase}/loss float32 (n_folds, n_epochs)
/{phase}/logits float32 (n_folds, n_epochs, n_samples, n_heads, n_classes)
/{phase}/head_names str (n_heads,)
/{phase}/y_true int64 | float64 (n_samples,) — float64 for regression targets, int64 otherwise
/{phase}/entity_id_{k} int64|str (n_samples,) — one dataset per id component
/{phase}/split str (n_folds, n_samples)
/{phase}/loss float32 (n_folds, n_epochs)
HDF5 layout — FeatureStore
---------------------------
/{phase}/{head_name} float32 (n_folds, n_epochs, n_samples, embedding_dim)
/{phase}/y_true int64 (n_samples,)
/{phase}/entity_id_{k} int64|str (n_samples,)
/{phase}/split str (n_folds, n_samples)
/{phase}/{head_name} float32 (n_folds, n_epochs, n_samples, embedding_dim)
/{phase}/y_true int64 | float64 (n_samples,) — float64 for regression targets, int64 otherwise
/{phase}/entity_id_{k} int64|str (n_samples,)
/{phase}/split str (n_folds, n_samples)
"""
from __future__ import annotations
@@ -35,6 +35,18 @@ except ImportError as e:
_STR_DT = h5py.string_dtype()
def _coerce_y_true(y_true) -> np.ndarray:
"""Coerce y_true to int64 for integer-typed input, float64 otherwise.
Forcing int64 unconditionally would silently round regression targets
(e.g. VF_MD), so we honour float input by storing as float64.
"""
arr = np.asarray(y_true)
if np.issubdtype(arr.dtype, np.floating):
return arr.astype(np.float64)
return arr.astype(np.int64)
# ---------------------------------------------------------------------------
# Internal phase buffer
# ---------------------------------------------------------------------------
@@ -52,7 +64,7 @@ class _PhaseBuffer:
n_s = len(entity_ids)
n_h = len(head_names)
self.entity_ids = list(entity_ids)
self.y_true = np.asarray(y_true, dtype=np.int64)
self.y_true = _coerce_y_true(y_true)
self.head_names = list(head_names)
self.n_epochs = n_epochs
self.logits = np.full((n_folds, n_epochs, n_s, n_h, n_classes), np.nan, dtype=np.float32)
@@ -74,7 +86,7 @@ class _FeaturePhaseBuffer:
n_folds: int,
):
self.entity_ids = list(entity_ids)
self.y_true = np.asarray(y_true, dtype=np.int64)
self.y_true = _coerce_y_true(y_true)
self.split = np.full((n_folds, len(entity_ids)), "", dtype=object)
self._sid = {str(eid): i for i, eid in enumerate(entity_ids)}
# head_name → (buffer array, n_epochs)
@@ -148,7 +160,7 @@ class PredictionStore:
"""Register a training phase before recording begins."""
self._phases[phase] = _PhaseBuffer(
entity_ids=list(entity_ids),
y_true=np.asarray(y_true, dtype=np.int64),
y_true=_coerce_y_true(y_true),
head_names=list(head_names),
n_epochs=n_epochs,
n_folds=self.n_folds,
@@ -306,7 +318,7 @@ class FeatureStore:
) -> None:
self._phases[phase] = _FeaturePhaseBuffer(
entity_ids=list(entity_ids),
y_true=np.asarray(y_true, dtype=np.int64),
y_true=_coerce_y_true(y_true),
n_folds=self.n_folds,
)
+171
View File
@@ -650,6 +650,177 @@ def build_geometry_loader(source: str, **kwargs):
raise NotImplementedError(f"build_geometry_loader: source={source!r} not implemented")
class _GTContourBboxLoader:
"""Disc bounding-box loader from PAPILA expert disc contours.
Computes a square bbox centred on the disc, expanded by ``margin`` ×
max(disc_w, disc_h). Returned bboxes are in original-image pixel coords
and may extend past image bounds (clip at crop time).
"""
def __init__(
self,
contour_dir: str | Path,
*,
margin: float = 2.5,
expert: int = 1,
) -> None:
self._contour_dir = Path(contour_dir)
self._margin = float(margin)
self._expert = int(expert)
self._cache: dict[tuple, tuple[int, int, int, int] | None] = {}
def reset_cache(self) -> None:
self._cache.clear()
def precompute(self, samples: Iterable[Tuple[int, str, Path]]) -> None:
for pid, eye, _ in list(samples):
key = (int(pid), str(eye))
if key in self._cache:
continue
self._cache[key] = self._compute_bbox(*key)
def bbox_for(self, pid, eye) -> tuple[int, int, int, int] | None:
key = (int(pid), str(eye))
if key not in self._cache:
self._cache[key] = self._compute_bbox(*key)
return self._cache[key]
def _compute_bbox(self, pid: int, eye: str) -> tuple[int, int, int, int] | None:
path = self._contour_dir / f"RET{pid:03d}{eye}_disc_exp{self._expert}.txt"
if not path.exists():
return None
try:
arr = np.loadtxt(str(path), dtype=np.float32)
except Exception:
return None
if arr.ndim == 1:
arr = arr.reshape(-1, 2)
if arr.shape[0] < 3:
return None
return _expand_bbox_from_points(arr[:, 0], arr[:, 1], self._margin)
class _UNetBboxLoader:
"""Disc bounding-box loader from U-Net predicted masks.
Reuses _PapilaUNetMaskPipeline for per-fold fine-tune + inference.
"""
def __init__(
self,
weights_path: str | Path,
*,
contour_dir: str | Path,
margin: float = 2.5,
unet_size: int = 512,
normalize: str = "per_image",
threshold: float = 0.5,
finetune_epochs: int = 0,
finetune_lr: float = 1e-5,
finetune_batch_size: int = 4,
device: str | None = None,
) -> None:
self._margin = float(margin)
self._unet_size = int(unet_size)
self._pipeline = _PapilaUNetMaskPipeline(
weights_path,
contour_dir=contour_dir,
unet_size=unet_size,
normalize=normalize,
threshold=threshold,
finetune_epochs=finetune_epochs,
finetune_lr=finetune_lr,
finetune_batch_size=finetune_batch_size,
device=device,
)
self._cache: dict[tuple, tuple[int, int, int, int] | None] = {}
def reset_cache(self) -> None:
self._cache.clear()
def reset_weights(self) -> None:
self._pipeline.reset_weights()
def finetune(self, train_samples: list) -> None:
self._pipeline.finetune(train_samples)
def precompute(self, samples: Iterable[Tuple[int, str, Path]]) -> None:
samples = list(samples)
if not samples:
return
# Need original-image dims to rescale mask coords back; capture per sample.
orig_sizes: dict[tuple, tuple[int, int]] = {}
for pid, eye, image_path in samples:
try:
with Image.open(image_path) as im:
orig_sizes[(int(pid), str(eye))] = im.size # (w, h)
except Exception:
continue
masks = self._pipeline.predict(samples)
for key, (disc, _cup) in masks.items():
ow, oh = orig_sizes.get(key, (self._unet_size, self._unet_size))
self._cache[key] = _bbox_from_mask(disc, ow, oh, self._margin)
def bbox_for(self, pid, eye) -> tuple[int, int, int, int] | None:
return self._cache.get((int(pid), str(eye)))
def _expand_bbox_from_points(
xs: np.ndarray, ys: np.ndarray, margin: float
) -> tuple[int, int, int, int]:
"""Square bbox centred on disc centroid, half-side = margin × max(w,h) / 2."""
x0, x1 = float(xs.min()), float(xs.max())
y0, y1 = float(ys.min()), float(ys.max())
cx, cy = (x0 + x1) / 2.0, (y0 + y1) / 2.0
half = max(x1 - x0, y1 - y0) * float(margin) / 2.0
return (int(round(cx - half)), int(round(cy - half)),
int(round(cx + half)), int(round(cy + half)))
def _bbox_from_mask(
mask: np.ndarray, orig_w: int, orig_h: int, margin: float,
) -> tuple[int, int, int, int] | None:
"""Compute original-image bbox from a binary mask at mask resolution."""
ys, xs = np.where(mask > 0)
if len(xs) == 0:
return None
sx = float(orig_w) / float(mask.shape[1])
sy = float(orig_h) / float(mask.shape[0])
return _expand_bbox_from_points(xs * sx, ys * sy, margin)
def build_disc_bbox_loader(source: str, **kwargs):
"""Return a disc-bbox loader for the given source.
Parameters
----------
source : "gt" | "unet"
GT kwargs:
contour_dir (required), margin=2.5, expert=1
UNet kwargs:
weights_path (required), contour_dir (required),
margin=2.5, unet_size=512, normalize="per_image", threshold=0.5,
finetune_epochs=0, finetune_lr=1e-5, finetune_batch_size=4, device=None
"""
if source == "gt":
if "contour_dir" not in kwargs:
raise ValueError("build_disc_bbox_loader source='gt' requires contour_dir")
gt_keys = {"contour_dir", "margin", "expert"}
return _GTContourBboxLoader(**{k: v for k, v in kwargs.items() if k in gt_keys})
if source == "unet":
if "weights_path" not in kwargs:
raise ValueError("build_disc_bbox_loader source='unet' requires weights_path")
if "contour_dir" not in kwargs:
raise ValueError(
"build_disc_bbox_loader source='unet' requires contour_dir "
"(needed for per-fold fine-tuning, even if finetune_epochs=0)"
)
return _UNetBboxLoader(**kwargs)
raise NotImplementedError(f"build_disc_bbox_loader: source={source!r} not implemented")
def build_seg_map_loader(source: str, **kwargs):
"""Return the appropriate seg-map loader for the given source string.
+6
View File
@@ -480,6 +480,12 @@ class ImageDataView:
from v4.classes.profiles.fundus_images import build_seg_map_loader as _build
return _build(source, **self._resolve_paths(kwargs))
def build_disc_bbox_loader(self, source: str, **kwargs):
"""Return a disc bounding-box loader for crop-to-disc preprocessing."""
from v4.classes.profiles.fundus_images import build_disc_bbox_loader as _build
kwargs.setdefault("contour_dir", self._DEFAULT_CONTOUR_DIR)
return _build(source, **self._resolve_paths(kwargs))
# ── Optional explainability hooks ────────────────────────────────────────
#
# These methods are consumed by v4.classes.accessory.explainability via
+81 -70
View File
@@ -1,6 +1,7 @@
"""stages/fusion — fusion stage runner: trains a bridge + associated head stages."""
from __future__ import annotations
import contextlib
import importlib
from random import choice, random as _random
@@ -8,6 +9,20 @@ import numpy as np
import torch
import torch.nn.functional as F
def _amp_ctx(cfg: dict, device):
"""Autocast context for forward+loss when training.amp is enabled.
bf16 is the default dtype because its dynamic range matches fp32 and no
GradScaler is required. Falls back to a no-op context when amp is disabled
or the device is not CUDA/ROCm.
"""
train_cfg = cfg.get("training", {})
if not train_cfg.get("amp", False) or getattr(device, "type", None) != "cuda":
return contextlib.nullcontext()
dtype = getattr(torch, train_cfg.get("amp_dtype", "bfloat16"))
return torch.autocast(device_type="cuda", dtype=dtype)
from v4.classes.dataset import LoaderShell, to_label_tensor
from v4.classes.metrics import score_arrays, compute_extended_metrics, tune_binary_threshold
from v4.classes.stages.helpers import (
@@ -223,82 +238,78 @@ def run(
if y_t.numel() == 0:
continue
if is_bilateral:
side_embs = {
side: encode_embedding(src, batch, side, towers, stage_models, cfg_stages, device)
for side, src in inputs.items()
}
local_embs = {name: bridge(side_embs)}
else:
local_embs = {n: towers[n](batch[n].to(device)) for n in inputs
if n in batch and torch.is_tensor(batch[n])}
if len(local_embs) != len(inputs):
continue
local_embs[name] = bridge(list(local_embs[n] for n in inputs))
loss = None
logits = None
head_logits = {
hs["name"]: head_models[hs["name"]](local_embs[hs["input"]])
for hs in head_stage_cfgs
if hs["input"] in local_embs
}
if is_bilateral or phase == "fused_warmup":
logits = head_logits.get(primary_hs_cfg["name"])
chosen_head = head_models.get(primary_hs_cfg["name"])
elif phase == "tower_warmup" and bcd_head_cfgs:
losses = [
head_compute_loss(head_models[hs["name"]],
head_logits[hs["name"]], batch, y_t,
class_weights=cw)
for hs in bcd_head_cfgs if hs["name"] in head_logits
]
if not losses:
continue
loss = sum(losses) / len(losses)
if hasattr(bridge, "modify_loss"):
loss = bridge.modify_loss(loss)
opt.zero_grad(); loss.backward(); opt.step()
total_loss += loss.item() * len(y_t)
total_n += len(y_t)
continue
elif tower_loss_mode == "all_losses" and bcd_head_cfgs:
# All-losses (v3 phase 3 control): sum primary + every aux head
# loss every step. Effective LR is implicitly N× single-head BCD
# — matches v3 semantics so the comparison is apples-to-apples.
all_head_names = ([primary_hs_cfg["name"]]
+ [hs["name"] for hs in bcd_head_cfgs])
losses = [
head_compute_loss(head_models[n], head_logits[n], batch, y_t,
class_weights=cw)
for n in all_head_names if n in head_logits
]
if not losses:
continue
loss = sum(losses)
if hasattr(bridge, "modify_loss"):
loss = bridge.modify_loss(loss)
opt.zero_grad(); loss.backward(); opt.step()
total_loss += loss.item() * len(y_t)
total_n += len(y_t)
continue
else:
if bcd_head_cfgs and _random() < bcd_prob:
chosen_hs = choice(bcd_head_cfgs)
with _amp_ctx(cfg, device):
if is_bilateral:
side_embs = {
side: encode_embedding(src, batch, side, towers, stage_models, cfg_stages, device)
for side, src in inputs.items()
}
local_embs = {name: bridge(side_embs)}
else:
chosen_hs = primary_hs_cfg
logits = head_logits.get(chosen_hs["name"])
chosen_head = head_models.get(chosen_hs["name"])
local_embs = {n: towers[n](batch[n].to(device)) for n in inputs
if n in batch and torch.is_tensor(batch[n])}
if len(local_embs) != len(inputs):
continue
local_embs[name] = bridge(list(local_embs[n] for n in inputs))
if logits is None:
head_logits = {
hs["name"]: head_models[hs["name"]](local_embs[hs["input"]])
for hs in head_stage_cfgs
if hs["input"] in local_embs
}
if is_bilateral or phase == "fused_warmup":
logits = head_logits.get(primary_hs_cfg["name"])
chosen_head = head_models.get(primary_hs_cfg["name"])
if logits is not None:
loss = head_compute_loss(chosen_head, logits, batch, y_t, class_weights=cw)
elif phase == "tower_warmup" and bcd_head_cfgs:
losses = [
head_compute_loss(head_models[hs["name"]],
head_logits[hs["name"]], batch, y_t,
class_weights=cw)
for hs in bcd_head_cfgs if hs["name"] in head_logits
]
if losses:
loss = sum(losses) / len(losses)
elif tower_loss_mode == "all_losses" and bcd_head_cfgs:
# All-losses (v3 phase 3 control): sum primary + every aux head
# loss every step. Effective LR is implicitly N× single-head BCD
# — matches v3 semantics so the comparison is apples-to-apples.
all_head_names = ([primary_hs_cfg["name"]]
+ [hs["name"] for hs in bcd_head_cfgs])
losses = [
head_compute_loss(head_models[n], head_logits[n], batch, y_t,
class_weights=cw)
for n in all_head_names if n in head_logits
]
if losses:
loss = sum(losses)
else:
if bcd_head_cfgs and _random() < bcd_prob:
chosen_hs = choice(bcd_head_cfgs)
else:
chosen_hs = primary_hs_cfg
logits = head_logits.get(chosen_hs["name"])
chosen_head = head_models.get(chosen_hs["name"])
if logits is not None:
loss = head_compute_loss(chosen_head, logits, batch, y_t, class_weights=cw)
if loss is not None and hasattr(bridge, "modify_loss"):
loss = bridge.modify_loss(loss)
if loss is None:
continue
loss = head_compute_loss(chosen_head, logits, batch, y_t, class_weights=cw)
if hasattr(bridge, "modify_loss"):
loss = bridge.modify_loss(loss)
opt.zero_grad(); loss.backward(); opt.step()
if logits.dim() >= 2:
if logits is not None and logits.dim() >= 2:
total_correct += int((logits.argmax(1) == y_t).sum())
total_loss += loss.item() * len(y_t)
total_n += len(y_t)
total_loss += loss.item() * len(y_t)
total_n += len(y_t)
tr_loss = total_loss / total_n if total_n else nan
tr_acc = total_correct / total_n if total_n else nan
+72 -7
View File
@@ -71,6 +71,13 @@ class ImageEncoder(TowerBase):
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.
crop_source : if set, crop each input image to a square disc-region
bbox before the standard transform pipeline. Values:
"gt" (use GT contour file) | "unet" (use U-Net mask)
| None (disabled, full-image pipeline).
crop_kwargs : dict forwarded to image_data.build_disc_bbox_loader().
Common keys: margin (default 2.5), expert (GT only),
weights_path / finetune_epochs (U-Net only).
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.
@@ -89,6 +96,10 @@ class ImageEncoder(TowerBase):
se_pre_norm: bool = True,
augment: bool = True,
cache_transformed: bool = False,
crop_size: int | None = None,
resize_size: int | None = None,
crop_source: str | None = None,
crop_kwargs: dict | None = None,
geometry_source: str | None = None,
**geom_kwargs: Any,
):
@@ -98,17 +109,37 @@ class ImageEncoder(TowerBase):
self.backbone, self._base_dim, self._blocks = build_backbone(backbone, freeze_ratio)
self._cache_transformed = cache_transformed
tf_kw = dict(crop_size=crop_size, resize_size=resize_size)
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._precache_tf, self._post_train_tf = build_split_transforms(
backbone, augment=augment, **tf_kw)
_, self._post_eval_tf = build_split_transforms(
backbone, augment=False, **tf_kw)
self._tensor_cache: dict[tuple, torch.Tensor] = {}
else:
self.transform = build_backbone_transform(backbone, augment=augment)
self.eval_transform = build_eval_transform(backbone)
self.transform = build_backbone_transform(backbone, augment=augment, **tf_kw)
self.eval_transform = build_eval_transform(backbone, **tf_kw)
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._bbox_loader = None
if crop_source is not None:
if not hasattr(image_data, "build_disc_bbox_loader"):
raise TypeError(
f"ImageEncoder crop_source={crop_source!r} requires "
f"image_data to implement build_disc_bbox_loader(), "
f"but {type(image_data).__name__} does not."
)
self._bbox_loader = image_data.build_disc_bbox_loader(
crop_source, **(crop_kwargs or {}),
)
print(
f"[ImageEncoder] crop_source={crop_source!r} "
f"kwargs={crop_kwargs or {}}",
flush=True,
)
self._geom_loader = None
if geometry_source is not None:
if not hasattr(image_data, "build_geometry_loader"):
@@ -134,16 +165,32 @@ class ImageEncoder(TowerBase):
def _side_map(self) -> dict[str, str]:
return self.image_data.side_map
def _load_image(self, *ids):
"""Load image, optionally cropped to the disc-region bbox."""
pil = self.image_data.load_image(*ids)
if self._bbox_loader is None:
return pil
bbox = self._bbox_loader.bbox_for(*ids[:2])
if bbox is None:
return pil
w, h = pil.size
x0, y0, x1, y1 = bbox
x0, y0 = max(0, x0), max(0, y0)
x1, y1 = min(w, x1), min(h, y1)
if x1 <= x0 or y1 <= y0:
return pil
return pil.crop((x0, y0, x1, y1))
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))
cached = self._precache_tf(self._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)
img = self._load_image(*ids)
t = self.transform if self.training else self.eval_transform
return t(img)
@@ -154,6 +201,24 @@ class ImageEncoder(TowerBase):
data = context.require("data")
split = context.require("split")
# Disc-region bbox precomputation (must run before any image load/cache).
if self._bbox_loader is not None:
train_samples = data.collect_samples(split.train)
all_samples = train_samples + data.collect_samples(split.val)
if split.test is not None:
all_samples += data.collect_samples(split.test)
if hasattr(self._bbox_loader, "reset_cache"):
self._bbox_loader.reset_cache()
if hasattr(self._bbox_loader, "reset_weights"):
self._bbox_loader.reset_weights()
if hasattr(self._bbox_loader, "finetune"):
self._bbox_loader.finetune(train_samples)
self._bbox_loader.precompute(all_samples)
print(
f"[ImageEncoder] precomputed disc bboxes for {len(all_samples)} samples",
flush=True,
)
if self._cache_transformed:
self._tensor_cache.clear()
n = self._warm_tensor_cache(data, split)
@@ -200,7 +265,7 @@ class ImageEncoder(TowerBase):
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))
self._tensor_cache[key] = self._precache_tf(self._load_image(pid, eye))
seen.add(key)
return len(self._tensor_cache)
+18 -6
View File
@@ -359,21 +359,28 @@ def main():
if save_predictions and eval_stage_preds:
# Collect all unique entity_ids across val+test sets of all folds.
# Preserve the natural dtype of y so regression targets keep their
# fractional values (casting to int silently rounds VF_MD).
seen, all_ids, id_to_y = set(), [], {}
y_is_float = False
for fp in eval_stage_preds:
for eid, y in zip(fp["val_ids"], fp["val_y"]):
k = str(eid)
if k not in seen:
seen.add(k); all_ids.append(eid)
id_to_y[k] = int(y)
y_is_float = y_is_float or np.issubdtype(np.asarray(y).dtype, np.floating)
id_to_y[k] = float(y) if y_is_float else int(y)
if fp.get("test_ids"):
for eid, y in zip(fp["test_ids"], fp["test_y"]):
k = str(eid)
if k not in seen:
seen.add(k); all_ids.append(eid)
id_to_y[k] = int(y)
y_is_float = y_is_float or np.issubdtype(np.asarray(y).dtype, np.floating)
id_to_y[k] = float(y) if y_is_float else int(y)
y_true = np.array([id_to_y.get(str(e), -1) for e in all_ids], dtype=np.int64)
sentinel = float("nan") if y_is_float else -1
dtype = np.float64 if y_is_float else np.int64
y_true = np.array([id_to_y.get(str(e), sentinel) for e in all_ids], dtype=dtype)
store = PredictionStore(n_folds=len(eval_stage_preds), n_classes=num_classes)
store.register_phase(
phase=eval_stage,
@@ -402,19 +409,24 @@ def main():
for phase, phase_preds in all_phase_preds.items():
emb_dim = phase_preds[0]["val_z"].shape[-1]
seen, all_ids, id_to_y = set(), [], {}
y_is_float = False
for fp in phase_preds:
for eid, y in zip(fp["val_ids"], fp["val_y"]):
k = str(eid)
if k not in seen:
seen.add(k); all_ids.append(eid)
id_to_y[k] = int(y)
y_is_float = y_is_float or np.issubdtype(np.asarray(y).dtype, np.floating)
id_to_y[k] = float(y) if y_is_float else int(y)
if fp.get("test_ids"):
for eid, y in zip(fp["test_ids"], fp["test_y"]):
k = str(eid)
if k not in seen:
seen.add(k); all_ids.append(eid)
id_to_y[k] = int(y)
y_true = np.array([id_to_y.get(str(e), -1) for e in all_ids], dtype=np.int64)
y_is_float = y_is_float or np.issubdtype(np.asarray(y).dtype, np.floating)
id_to_y[k] = float(y) if y_is_float else int(y)
sentinel = float("nan") if y_is_float else -1
dtype = np.float64 if y_is_float else np.int64
y_true = np.array([id_to_y.get(str(e), sentinel) for e in all_ids], dtype=dtype)
fstore.register_phase(phase=phase, entity_ids=all_ids, y_true=y_true)
fstore.register_head(phase=phase, head=f"{phase}_embedding",
n_epochs=1, embedding_dim=emb_dim)