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:
@@ -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()
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
|
||||
@@ -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.
|
||||
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user