ryan/hypertower
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add new regression and ensemble experiment configurations for V2-M and OrthoBridge

Browse Source 
- Introduced multiple regression experiment configurations targeting vf_md, including:
  - cd_solo_reg_set.json: CD tower only regression setup.
  - img_solo_reg_set.json: Image tower only regression setup.
  - reg_head_epoch_sweep.json: Baseline regression sweeps at different epochs (50, 75, 100).
  - reg_head_set.json: Various regression setups including baseline and OrthoBridge configurations.
  - single_eye_reg.json: Single-eye regression setup for worst-eye aggregation analysis.

- Added ensemble configurations for OrthoBridge with different inner bridges:
  - ortho_alts_ensemble.json: Ensemble tests with ConcatBridge, PairwiseAdditiveBridge, and GatedAdditiveBridge.
  - ortho_alts_tritower.json: Tritower tests with the same inner bridges.

- Created V2-M specific configurations:
  - baseline_reg_nt50.json: Regression baseline with V2-M backbone.
  - geom_vec_gt.json and geom_vec_unet.json: Geometry vector injection experiments with V2-M.
  - single_l1_bridges.json: Single-eye ensemble experiments with various bridge types.
  - tritower_geom_gt.json: Tritower setup with GT contour-rasterized masks.

- Promoted existing experiments to higher repetitions for robustness.
This commit is contained in:
rpotter6298
2026-06-11 15:08:20 +02:00
parent 32a801a572
commit 280060db82
343 changed files with 8558 additions and 57747 deletions
Download Patch File Download Diff File Expand all files Collapse all files
v4/classes/accessory/backbones.py
+101 -11
View File
@@ -12,10 +12,11 @@ from torchvision import models, transforms
_REPO_ROOT = Path(__file__).resolve().parents[3]
REFUGELIKE_BACKBONE_PATH = _REPO_ROOT / "models/v2/refuge/refugelike_backbone.pt"
REFUGE_DENSENET_PATH = _REPO_ROOT / "models/refuge/classifier/refuge_densenet_backbone.pt"
REFUGE_EFFICIENT_B0_PATH = _REPO_ROOT / "models/refuge/classifier/refuge_efficient_b0_backbone.pt"
REFUGE_EFFICIENT_B7_PATH = _REPO_ROOT / "models/refuge/classifier/refuge_efficient_b7_backbone.pt"
REFUGELIKE_BACKBONE_PATH = _REPO_ROOT / "models/v2/refuge/refugelike_backbone.pt"
REFUGE_DENSENET_PATH = _REPO_ROOT / "models/refuge/classifier/refuge_densenet_backbone.pt"
REFUGE_EFFICIENTNET_B0_PATH = _REPO_ROOT / "models/refuge/classifier/refuge_efficientnet_b0_backbone.pt"
REFUGE_EFFICIENTNET_B7_PATH = _REPO_ROOT / "models/refuge/classifier/refuge_efficientnet_b7_backbone.pt"
REFUGE_EFFICIENTNET_V2_M_PATH = _REPO_ROOT / "models/refuge/classifier/refuge_efficientnet_v2_m_backbone.pt"
@dataclass(frozen=True)
@@ -109,6 +110,30 @@ BACKBONES: Dict[str, BackboneSpec] = {
strip=_strip_efficientnet,
blocks=_blocks_efficientnet,
),
"efficientnet_b7": BackboneSpec(
ctor=models.efficientnet_b7,
weights_default=models.EfficientNet_B7_Weights.DEFAULT,
strip=_strip_efficientnet,
blocks=_blocks_efficientnet,
),
"efficientnet_v2_s": BackboneSpec(
ctor=models.efficientnet_v2_s,
weights_default=models.EfficientNet_V2_S_Weights.DEFAULT,
strip=_strip_efficientnet,
blocks=_blocks_efficientnet,
),
"efficientnet_v2_m": BackboneSpec(
ctor=models.efficientnet_v2_m,
weights_default=models.EfficientNet_V2_M_Weights.DEFAULT,
strip=_strip_efficientnet,
blocks=_blocks_efficientnet,
),
"efficientnet_v2_l": BackboneSpec(
ctor=models.efficientnet_v2_l,
weights_default=models.EfficientNet_V2_L_Weights.DEFAULT,
strip=_strip_efficientnet,
blocks=_blocks_efficientnet,
),
"resnet18": BackboneSpec(
ctor=models.resnet18,
weights_default=models.ResNet18_Weights.DEFAULT,
@@ -157,31 +182,94 @@ BACKBONES: Dict[str, BackboneSpec] = {
strip=_strip_densenet,
blocks=_blocks_densenet,
),
"refuge_efficient_b0": BackboneSpec(
"refuge_efficientnet_b0": BackboneSpec(
ctor=models.efficientnet_b0,
weights_default=None,
strip=_strip_efficientnet,
blocks=_blocks_efficientnet,
),
"refuge_efficient_b7": BackboneSpec(
"refuge_efficientnet_b7": BackboneSpec(
ctor=models.efficientnet_b7,
weights_default=None,
strip=_strip_efficientnet,
blocks=_blocks_efficientnet,
),
"refuge_efficientnet_v2_m": BackboneSpec(
ctor=models.efficientnet_v2_m,
weights_default=None,
strip=_strip_efficientnet,
blocks=_blocks_efficientnet,
),
}
_TIMM_CONVNEXTV2_VARIANTS = (
"convnextv2_atto",
"convnextv2_femto",
"convnextv2_pico",
"convnextv2_nano",
"convnextv2_tiny",
"convnextv2_base",
"convnextv2_large",
"convnextv2_huge",
)
def _is_timm_backbone(name: str) -> bool:
return name in _TIMM_CONVNEXTV2_VARIANTS
def _build_timm_backbone(name: str, freeze_ratio: float) -> tuple[nn.Module, int, list]:
"""Build a timm-sourced backbone (ConvNeXt-V2 family).
Uses num_classes=0 + global_pool="avg" so the model returns a pooled
(B, num_features) tensor directly — no head to strip. Freezable blocks
are [stem, stage_0, stage_1, stage_2, stage_3].
"""
try:
import timm
except ImportError as e:
raise ImportError(
f"backbone {name!r} requires the `timm` package "
"(pip install timm)"
) from e
m = timm.create_model(name, pretrained=True, num_classes=0, global_pool="avg")
out_dim = int(getattr(m, "num_features", 0))
if not out_dim:
raise RuntimeError(f"timm model {name!r} did not expose num_features")
blocks: list[nn.Module] = []
if hasattr(m, "stem"):
blocks.append(m.stem)
if hasattr(m, "stages"):
blocks.extend(list(m.stages))
if not blocks:
raise RuntimeError(
f"timm model {name!r} did not expose .stem / .stages — "
"extend _build_timm_backbone to support this architecture."
)
fr = max(0.0, min(1.0, float(freeze_ratio)))
n_freeze = int(math.floor(len(blocks) * fr))
for b in blocks[:n_freeze]:
for p in b.parameters():
p.requires_grad = False
return m, out_dim, blocks
def list_names() -> List[str]:
return list(BACKBONES.keys())
return list(BACKBONES.keys()) + list(_TIMM_CONVNEXTV2_VARIANTS)
def load_backbone_weights(key: str, model: nn.Module) -> None:
paths = {
"refugelike": REFUGELIKE_BACKBONE_PATH,
"refuge_densenet": REFUGE_DENSENET_PATH,
"refuge_efficient_b0": REFUGE_EFFICIENT_B0_PATH,
"refuge_efficient_b7": REFUGE_EFFICIENT_B7_PATH,
"refugelike": REFUGELIKE_BACKBONE_PATH,
"refuge_densenet": REFUGE_DENSENET_PATH,
"refuge_efficientnet_b0": REFUGE_EFFICIENTNET_B0_PATH,
"refuge_efficientnet_b7": REFUGE_EFFICIENTNET_B7_PATH,
"refuge_efficientnet_v2_m": REFUGE_EFFICIENTNET_V2_M_PATH,
}
path = paths.get(key)
if path is None:
@@ -202,6 +290,8 @@ def build_backbone(name: str, freeze_ratio: float = 0.0) -> tuple[nn.Module, int
freezable units — callers use it to dynamically adjust freeze_ratio later.
"""
key = (name or "").lower()
if _is_timm_backbone(key):
return _build_timm_backbone(key, freeze_ratio)
if key not in BACKBONES:
raise ValueError(f"Unknown backbone '{name}'. Available: {list_names()}")
v4/classes/accessory/explainability.py
+304
View File
@@ -0,0 +1,304 @@
"""Generic explainability primitives — GradCAM + region-of-interest analysis
+ column-shuffling permutation importance.
This module is dataset-agnostic. All PAPILA-specific knowledge (where the disc
contour lives, eye orientation conventions, clinical feature names, etc.) is
supplied by the profile through optional methods consumed via getattr — see
v4.classes.profiles.v4papila.ImageDataView and ClinicalDataView for examples.
Public surface
--------------
GradCAM — hooks-based class-activation maps via fwd/bwd hooks
overlay_gradcam — render a heatmap blended onto a PIL image
cam_region_stats — quantify CAM attention inside/outside a binary ROI
cam_region_patch — crop a square patch centered on the ROI centroid
permutation_importance — per-feature drop-in-score under column shuffling
"""
from __future__ import annotations
from typing import Callable, Sequence
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from matplotlib import cm
from PIL import Image
# ─────────────────────────────────────────────────────────────────────────────
# GradCAM
# ─────────────────────────────────────────────────────────────────────────────
class GradCAM:
"""Minimal GradCAM via forward/backward hooks.
Architecture-agnostic: the caller supplies a forward_fn that returns a
(1, num_classes) logits tensor, so this class doesn't need to know whether
the model takes (img,), (img, meta), or composed tower/bridge/head inputs.
Usage
-----
cam_module = GradCAM(target_layer)
cam, pred = cam_module.compute(
forward_fn=lambda: head(bridge([tower_img(x), tower_cd(c)])),
output_shape=(224, 224),
)
cam_module.remove()
"""
def __init__(self, target_layer: nn.Module) -> None:
self._acts: torch.Tensor | None = None
self._grads: torch.Tensor | None = None
self._h1 = target_layer.register_forward_hook(self._save_acts)
self._h2 = target_layer.register_full_backward_hook(self._save_grads)
def _save_acts(self, _module, _inputs, output):
self._acts = output.detach()
def _save_grads(self, _module, _grad_in, grad_out):
self._grads = grad_out[0].detach()
def compute(
self,
forward_fn: Callable[[], torch.Tensor],
*,
output_shape: tuple[int, int],
target_class: int | None = None,
) -> tuple[np.ndarray, int]:
"""Run forward via forward_fn, backward against target_class, return CAM.
Parameters
----------
forward_fn
Zero-arg callable returning a (1, C) logits tensor. Caller is
responsible for placing model in eval mode and zeroing grads.
output_shape
(H, W) at which the CAM is bilinearly upsampled. Typically the
eval-crop image size, e.g. (224, 224).
target_class
Class index to backprop against. If None, uses the predicted class.
Returns
-------
(cam, pred)
cam : float32 ndarray of shape output_shape, normalised to [0, 1]
pred : int, the predicted class index
"""
with torch.enable_grad():
out = forward_fn()
pred = int(out.argmax(1).item())
tc = pred if target_class is None else target_class
out[0, tc].backward()
assert self._grads is not None and self._acts is not None, (
"GradCAM hooks did not fire — check that target_layer is in the "
"forward path of forward_fn()."
)
weights = self._grads.mean(dim=(2, 3), keepdim=True)
cam = F.relu((weights * self._acts).sum(dim=1, keepdim=True))
cam = F.interpolate(cam, output_shape, mode="bilinear", align_corners=False)
cam_np = cam.squeeze().cpu().numpy()
lo, hi = cam_np.min(), cam_np.max()
return ((cam_np - lo) / (hi - lo + 1e-8)).astype(np.float32), pred
def remove(self) -> None:
"""Detach the forward/backward hooks. Call when done with this instance."""
self._h1.remove()
self._h2.remove()
# ─────────────────────────────────────────────────────────────────────────────
# Visualisation
# ─────────────────────────────────────────────────────────────────────────────
def overlay_gradcam(pil: Image.Image, cam: np.ndarray, alpha: float = 0.45) -> Image.Image:
"""Blend a jet-colored heatmap of `cam` onto `pil` at alpha blending weight."""
cam_u8 = (np.clip(cam, 0.0, 1.0) * 255).astype(np.uint8)
cam_r = np.array(Image.fromarray(cam_u8).resize(pil.size, Image.BILINEAR)) / 255.0
colored = (cm.jet(cam_r)[:, :, :3] * 255).astype(np.uint8)
return Image.blend(pil.convert("RGB"), Image.fromarray(colored), alpha)
# ─────────────────────────────────────────────────────────────────────────────
# Region-of-interest analysis (generic — no PAPILA knowledge)
# ─────────────────────────────────────────────────────────────────────────────
def cam_region_stats(cam: np.ndarray, roi_mask: np.ndarray) -> dict[str, float]:
"""Quantify CAM attention restricted to / outside a binary ROI mask.
The CAM and the mask must be at the same resolution.
Returns
-------
dict with:
inside_mean : mean CAM intensity inside the ROI
outside_mean : mean CAM intensity outside the ROI
inside_fraction : sum(CAM inside) / sum(CAM total); fraction of total
attention that falls inside the ROI
roi_area_frac : ROI area / total area; how much of the image the ROI
covers — useful for normalising inside_fraction
"""
if cam.shape != roi_mask.shape:
raise ValueError(
f"cam shape {cam.shape} != roi_mask shape {roi_mask.shape}"
)
mask = roi_mask.astype(bool)
cam_sum = float(cam.sum())
inside_sum = float(cam[mask].sum()) if mask.any() else 0.0
outside_sum = cam_sum - inside_sum
n_inside = int(mask.sum())
n_outside = int(mask.size - n_inside)
return {
"inside_mean": inside_sum / max(1, n_inside),
"outside_mean": outside_sum / max(1, n_outside),
"inside_fraction": inside_sum / cam_sum if cam_sum > 0 else 0.0,
"roi_area_frac": n_inside / mask.size,
}
def cam_region_patch(
cam: np.ndarray,
roi_mask: np.ndarray,
*,
span: float = 5.0,
patch_size: int = 96,
) -> tuple[np.ndarray | None, float | None]:
"""Crop a square patch centered on the ROI centroid.
The crop side is `span × roi_radius` (in CAM pixels), where roi_radius is
derived from the mask area assuming a circular ROI. The output is resized
to (patch_size, patch_size).
Returns
-------
(patch, roi_radius_in_patch_pixels)
patch : float32 ndarray of shape (patch_size, patch_size), values in
[0, 1]. Padded with zeros if the centered crop window extends
outside the CAM.
roi_radius_in_patch_pixels : float, the ROI radius re-expressed in
output-patch pixels (useful for drawing ROI circles in figures).
Returns (None, None) if the mask is empty.
"""
mask = roi_mask.astype(bool)
if not mask.any():
return None, None
ys, xs = np.where(mask)
cy, cx = ys.mean(), xs.mean()
roi_r = float(np.sqrt(mask.sum() / np.pi))
half = max(1, int(round(span * roi_r / 2)))
h, w = cam.shape
y0, y1 = int(round(cy)) - half, int(round(cy)) + half
x0, x1 = int(round(cx)) - half, int(round(cx)) + half
pt = max(0, -y0)
pb = max(0, y1 - h)
pl = max(0, -x0)
pr = max(0, x1 - w)
cam_pad = np.pad(cam, ((pt, pb), (pl, pr)), constant_values=0.0)
window = cam_pad[y0 + pt: y1 + pt, x0 + pl: x1 + pl]
patch = np.array(
Image.fromarray((np.clip(window, 0, 1) * 255).astype(np.uint8))
.resize((patch_size, patch_size), Image.BILINEAR)
) / 255.0
roi_r_out = patch_size * roi_r / (2 * half)
return patch.astype(np.float32), roi_r_out
# ─────────────────────────────────────────────────────────────────────────────
# Permutation importance (column-shuffling)
# ─────────────────────────────────────────────────────────────────────────────
def permutation_importance(
*,
score_fn: Callable[[np.ndarray], float],
X: np.ndarray,
groups: Sequence[Sequence[int]] | None = None,
n_permutations: int = 30,
seed: int = 0,
feature_names: Sequence[str] | None = None,
) -> dict:
"""Per-feature drop-in-score from column shuffling.
Generic and feature-agnostic — the caller wires the model (or any
arbitrary scoring pipeline) into ``score_fn``. Higher score = better;
drop = baseline_score − score(X_with_column_permuted).
Parameters
----------
score_fn
Callable taking the input matrix ``X`` (shape (N, F)) and returning
a scalar metric. Caller is responsible for any constants the metric
depends on (e.g. labels, held-out image embeddings).
X
Float ndarray of shape (N, F). The base matrix to permute.
groups
Optional list of column-index groups. If provided, each group is
permuted together (useful when several model dims encode one
conceptual feature, e.g. one-hot encodings of a categorical column).
Defaults to one group per column.
n_permutations
Number of independent shuffles per group; mean and std of the drop
are reported across these.
seed
Seed for the shuffling RNG.
feature_names
Optional human-readable names, one per group. If not provided, names
default to f"group_{i}" or f"col_{j}" when groups is None.
Returns
-------
dict with keys:
baseline : float, the score on unmodified X
mean_drop : ndarray of shape (G,), mean (baseline − shuffled) per group
std_drop : ndarray of shape (G,), std of per-permutation drops
feature_names : list[str] of length G
groups : list[list[int]] of resolved column groups
"""
if X.ndim != 2:
raise ValueError(f"X must be 2-D (N, F); got shape {X.shape}")
n, f_dim = X.shape
if groups is None:
resolved_groups: list[list[int]] = [[j] for j in range(f_dim)]
else:
resolved_groups = [list(g) for g in groups]
all_cols = [c for g in resolved_groups for c in g]
if any(c < 0 or c >= f_dim for c in all_cols):
raise ValueError(
f"groups contain column indices outside [0, {f_dim})"
)
n_groups = len(resolved_groups)
if feature_names is None:
names = (
[f"col_{g[0]}" for g in resolved_groups] if groups is None
else [f"group_{i}" for i in range(n_groups)]
)
else:
names = list(feature_names)
if len(names) != n_groups:
raise ValueError(
f"feature_names has length {len(names)} but there are "
f"{n_groups} groups"
)
rng = np.random.default_rng(seed)
baseline = float(score_fn(X))
drops = np.empty((n_groups, n_permutations), dtype=float)
for gi, cols in enumerate(resolved_groups):
for p in range(n_permutations):
X_perm = X.copy()
perm = rng.permutation(n)
X_perm[:, cols] = X_perm[perm][:, cols]
drops[gi, p] = baseline - float(score_fn(X_perm))
return {
"baseline": baseline,
"mean_drop": drops.mean(axis=1),
"std_drop": drops.std(axis=1),
"feature_names": names,
"groups": resolved_groups,
}
v4/classes/accessory/transforms.py
+6 -1
View File
@@ -6,7 +6,7 @@ from typing import Tuple
from torchvision import transforms
from v4.classes.accessory.backbones import BACKBONES
from v4.classes.accessory.backbones import BACKBONES, _is_timm_backbone
IMAGENET_MEAN: Tuple[float, float, float] = (0.485, 0.456, 0.406)
IMAGENET_STD: Tuple[float, float, float] = (0.229, 0.224, 0.225)
@@ -79,6 +79,11 @@ class ImageTransformConfig:
def backbone_transform_config(backbone_name: str, augment: bool = True) -> ImageTransformConfig:
"""Build an ImageTransformConfig using the backbone's default normalisation stats."""
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]
v4/classes/dataset.py
+13 -1
View File
@@ -208,10 +208,22 @@ class HTDataset(Dataset):
def __getitem__(self, idx: int) -> dict[str, Any]:
entry = self.entries[idx]
sample = {
sample: dict[str, Any] = {
"label": torch.tensor(entry.label, dtype=torch.long),
"entity_id": entry.entity_id,
}
# Any per-entry auxiliary fields the profile attached via meta
# (e.g. vf_md for regression heads) flow into the batch dict
# alongside `label` and tower inputs.
for k, v in entry.meta.items():
if k in sample:
continue # don't overwrite primary fields
if isinstance(v, torch.Tensor):
sample[k] = v
elif isinstance(v, (int, float)):
sample[k] = torch.tensor(v, dtype=torch.float32)
else:
sample[k] = v
for name, tower in self.towers.items():
sample[name] = tower.get_sample(entry)
return sample
v4/classes/heads/regression.py
+165
View File
@@ -0,0 +1,165 @@
"""regression — RegressionHead for continuous-target tasks.
Self-contained head: implements the four opt-in hooks expected by the v4
stage runners (`target_key`, `compute_loss`, `to_probs`, `score`), so no
runner-side classification logic interferes.
Usage in a stage config:
{
"name": "nt_head",
"type": "head",
"input": "nt",
"train_with": "nt",
"module": "v4.classes.heads.regression",
"class": "RegressionHead",
"args": {
"target_key": "vf_md",
"dropout": 0.3,
"loss": "mse"
}
}
"""
from __future__ import annotations
from typing import Any
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
_LOSS_FNS = {
"mse": F.mse_loss,
"l1": F.l1_loss,
"huber": F.huber_loss,
}
class RegressionHead(nn.Module):
"""Minimal regression head: ReLU → Dropout → Linear(in_dim → 1).
Parameters
----------
in_dim : input embedding dimension (set by the stage runner)
num_classes : ignored — present only so the runner's `h_cls(in_dim, num_classes, **args)`
invocation keeps working. The head always outputs 1 scalar.
dropout : dropout fraction before the final linear
target_key : which batch field to read as the ground-truth value
(must match a field populated by the dataset; default 'vf_md')
loss : 'mse' | 'l1' | 'huber'
"""
def __init__(
self,
in_dim: int,
num_classes: int = 1,
dropout: float = 0.3,
target_key: str = "vf_md",
loss: str = "mse",
):
super().__init__()
del num_classes # explicitly ignored; always single-output
self.target_key = target_key
if loss not in _LOSS_FNS:
raise ValueError(f"unknown loss {loss!r}; choose from {list(_LOSS_FNS)}")
self._loss_name = loss
self._loss_fn = _LOSS_FNS[loss]
self.net = nn.Sequential(
nn.ReLU(),
nn.Dropout(dropout),
nn.Linear(in_dim, 1),
)
# ── nn.Module ──────────────────────────────────────────────────────────────
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.net(x).squeeze(-1) # (B,) — per-sample predicted scalar
# ── opt-in hooks consumed by the stage runners ─────────────────────────────
def compute_loss(self, logits: torch.Tensor, batch: dict) -> torch.Tensor:
"""Run regression loss against this head's target field in the batch."""
target = batch[self.target_key]
if not torch.is_tensor(target):
target = torch.as_tensor(target)
target = target.float().to(logits.device)
# logits is (B,), target is (B,) — match shapes for any loss fn
return self._loss_fn(logits, target)
def to_probs(self, logits: torch.Tensor) -> np.ndarray:
"""For regression there is no 'probability' — return raw predictions."""
return logits.detach().cpu().numpy()
def score(self, y_true: np.ndarray, predictions: np.ndarray) -> dict:
"""Compute regression metrics. Returns dict with 'primary' = neg-MSE.
We report negative MSE as 'primary' so that 'higher is better' matches
the convention of AUC — easier for monitoring code that expects to
maximise the primary metric. MSE itself is also stored separately.
"""
y = np.asarray(y_true, dtype=np.float32).ravel()
pred = np.asarray(predictions, dtype=np.float32).ravel()
n = y.size
if n == 0:
return {"primary": float("nan"),
"primary_name": "neg_mse",
"neg_mse": float("nan"),
"mse": float("nan"),
"mae": float("nan"),
"r2": float("nan"),
"spearman": float("nan"),
"n": 0}
diff = pred - y
mse = float(np.mean(diff ** 2))
mae = float(np.mean(np.abs(diff)))
# R²: 1 - SSres / SStot; falls back to NaN if SStot = 0
ss_res = float(np.sum(diff ** 2))
ss_tot = float(np.sum((y - y.mean()) ** 2))
r2 = float(1.0 - ss_res / ss_tot) if ss_tot > 1e-12 else float("nan")
# Spearman rank correlation (no scipy dependency — manual rank-Pearson)
spearman = _spearmanr(y, pred)
return {
"primary": -mse, # neg-MSE so 'higher is better'
"primary_name": "neg_mse",
"neg_mse": -mse, # mirror under its primary_name (matches AUC pattern)
"mse": mse,
"mae": mae,
"r2": r2,
"spearman": spearman,
"n": int(n),
}
def _spearmanr(a: np.ndarray, b: np.ndarray) -> float:
"""Simple Spearman correlation — Pearson on ranks; handles ties via average ranks."""
if a.size < 2:
return float("nan")
ra = _rankdata(a)
rb = _rankdata(b)
ra -= ra.mean(); rb -= rb.mean()
den = float(np.sqrt((ra ** 2).sum() * (rb ** 2).sum()))
if den < 1e-12:
return float("nan")
return float((ra * rb).sum() / den)
def _rankdata(x: np.ndarray) -> np.ndarray:
"""Assign average ranks; equivalent to scipy.stats.rankdata."""
order = np.argsort(x, kind="mergesort")
ranks = np.empty_like(order, dtype=np.float64)
ranks[order] = np.arange(1, len(x) + 1)
# Average ties
sorted_x = x[order]
i = 0
while i < len(x):
j = i
while j + 1 < len(x) and sorted_x[j + 1] == sorted_x[i]:
j += 1
if j > i:
avg = (ranks[order[i:j + 1]].mean())
ranks[order[i:j + 1]] = avg
i = j + 1
return ranks
v4/classes/profiles/fundus_images.py
+2 -1
View File
@@ -667,7 +667,8 @@ def build_seg_map_loader(source: str, **kwargs):
if source == "gt":
if "contour_dir" not in kwargs:
raise ValueError("build_seg_map_loader source='gt' requires contour_dir")
return GTSegMapLoader(**kwargs)
gt_keys = {"contour_dir", "channels", "mask_size", "target_size", "crop_to_disc"}
return GTSegMapLoader(**{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_seg_map_loader source='unet' requires weights_path")
v4/classes/profiles/v4papila.py
+260 -7
View File
@@ -155,7 +155,12 @@ def _apply_iop_and_drop_md(
df["IOP_corr"] = [
_correct_iop(r, p) for r, p in zip(df["IOP_raw"].values, pachy.values)
]
drop = [c for c in ("Pneumatic", "Perkins", "VF_MD") if c in df.columns]
# Drop the raw IOP source columns (Pneumatic, Perkins) since their info
# is already absorbed into IOP_corr. VF_MD is intentionally NOT dropped
# here — it stays available as a target for auxiliary/regression tasks
# (PapilaBundle exposes it via .vf_md_target / .get_vf_md()). It must
# however be added to `exclude_cols` so it is never used as a feature.
drop = [c for c in ("Pneumatic", "Perkins") if c in df.columns]
if drop_raw:
drop.append("IOP_raw")
if drop:
@@ -248,6 +253,39 @@ class ClinicalDataView:
n_cat = sum(len(m) for m in self.cat_maps.values())
return n_scalar + n_cat + n_scalar # scalars + one-hots + missing flags
# ── Optional explainability hooks ────────────────────────────────────────
#
# feature_names + feature_groups describe how the encoded feature_dim
# vector maps back to the original column space. They're consumed by
# v4.classes.accessory.explainability.permutation_importance via the
# caller (F8); profiles without these methods skip per-column importance.
@cached_property
def feature_names(self) -> list[str]:
"""Human-readable name for each original column (used as group label)."""
return list(self.scalar_cols) + list(self.cat_cols)
@cached_property
def feature_groups(self) -> list[list[int]]:
"""Encoded-vector indices grouped by original column.
Each entry is the set of model-input dimensions that encode one
conceptual feature: a scalar groups its value + missing-flag dim,
a categorical groups all of its one-hot dims.
Order matches ``feature_names``.
"""
n_scalar = len(self.scalar_cols)
n_cat_total = sum(len(m) for m in self.cat_maps.values())
groups: list[list[int]] = []
for i in range(n_scalar):
groups.append([i, n_scalar + n_cat_total + i])
offset = n_scalar
for col in self.cat_cols:
k = len(self.cat_maps[col])
groups.append(list(range(offset, offset + k)))
offset += k
return groups
def vectorize_entity(self, *ids) -> np.ndarray:
"""Return the feature vector for an entity identified by positional ids.
@@ -442,6 +480,104 @@ class ImageDataView:
from v4.classes.profiles.fundus_images import build_seg_map_loader as _build
return _build(source, **self._resolve_paths(kwargs))
# ── Optional explainability hooks ────────────────────────────────────────
#
# These methods are consumed by v4.classes.accessory.explainability via
# getattr — they're optional on the data view, so a different profile
# without disc annotations can simply not define them and GradCAM/overlay
# will still work (region-of-interest analysis is skipped).
#
# Eval-crop convention: PAPILA training resizes the short side to 256
# then center-crops 224×224. The methods below mirror that so the disc
# mask and the display image stay aligned with what the image tower sees.
_EVAL_RESIZE = 256
_EVAL_CROP = 224
_DEFAULT_CONTOUR_DIR = "Papila/ExpertsSegmentations/Contours"
def eval_image_pil(self, *ids) -> Image.Image:
"""Load image and apply the 256-resize + 224-center-crop used at eval time."""
from torchvision.transforms import functional as TF
from torchvision.transforms import InterpolationMode
pil = self.load_image(*ids).convert("RGB")
pil = TF.resize(pil, self._EVAL_RESIZE, interpolation=InterpolationMode.BILINEAR)
return TF.center_crop(pil, [self._EVAL_CROP, self._EVAL_CROP])
def build_roi_mask(
self,
*ids,
target_h: int,
target_w: int,
expert: int = 1,
contour_dir: str | None = None,
) -> np.ndarray | None:
"""Rasterize the expert disc contour for (pid, eye), aligned to the eval crop.
Returns a binary ndarray of shape (target_h, target_w), or None if no
contour file exists for this eye / the file is malformed.
Alignment: the contour polygon is rasterized at the original image
resolution, then resized + center-cropped to match the same eval
transform the image tower uses, then resized to (target_h, target_w).
"""
from PIL import ImageDraw
from torchvision.transforms import functional as TF
from torchvision.transforms import InterpolationMode
if len(ids) < 2:
raise TypeError(
"build_roi_mask requires (patient_id, eye) — got %r" % (ids,)
)
pid, eye = int(ids[0]), str(ids[1])
repo_root = Path(__file__).resolve().parents[3]
dir_ = Path(contour_dir or self._DEFAULT_CONTOUR_DIR)
if not dir_.is_absolute():
dir_ = repo_root / dir_
contour_path = dir_ / f"RET{pid:03d}{eye}_disc_exp{expert}.txt"
if not contour_path.exists():
return None
try:
arr = np.loadtxt(str(contour_path), dtype=np.float32)
except Exception:
return None
if arr.ndim == 1:
arr = arr.reshape(-1, 2)
if arr.shape[0] < 3:
return None
# Original image size for polygon canvas
try:
with Image.open(self.get_image_path(*ids)) as im:
orig_w, orig_h = im.size
except Exception:
return None
canvas = Image.new("L", (orig_w, orig_h), 0)
ImageDraw.Draw(canvas).polygon([tuple(pt) for pt in arr[:, :2]], fill=1)
# Match the eval transform exactly (resize → center-crop)
canvas = TF.resize(canvas, self._EVAL_RESIZE, interpolation=InterpolationMode.NEAREST)
canvas = TF.center_crop(canvas, [self._EVAL_CROP, self._EVAL_CROP])
if (target_h, target_w) != (self._EVAL_CROP, self._EVAL_CROP):
canvas = canvas.resize((target_w, target_h), Image.NEAREST)
return np.array(canvas, dtype=bool)
def orient_for_display(self, image, side: str):
"""Mirror OS to align its nasal–temporal axis with the OD convention.
Accepts a PIL.Image or a numpy array. OD is passed through unchanged.
"""
if side != self.SIDE_B: # SIDE_A == OD, unchanged
return image
if isinstance(image, Image.Image):
from PIL import ImageOps
return ImageOps.mirror(image)
if isinstance(image, np.ndarray):
return np.ascontiguousarray(np.fliplr(image))
raise TypeError(f"orient_for_display: unsupported type {type(image).__name__}")
# ---------------------------------------------------------------------------
# PapilaBundle — the v4 DataBundle returned by build_data
@@ -457,10 +593,11 @@ class PapilaBundle:
def __init__(
self,
bundle: DataBundle,
image_dir: str,
preprocessor: Callable | None = None,
image_cache: CachedImageLoader | None = None,
bundle: DataBundle,
image_dir: str,
preprocessor: Callable | None = None,
image_cache: CachedImageLoader | None = None,
vf_md_targets: dict[tuple, float] | None = None,
):
self._bundle = bundle
self._image_dir = image_dir
@@ -485,6 +622,25 @@ class PapilaBundle:
image_cache=image_cache,
)
# ── VF_MD target table (for regression / auxiliary heads) ────────────
# Provided by build_data — captured from raw frames before the master
# df dropped VF_MD via exclude_cols. Keys are (pid, eyeID), values are
# raw floats (NaN where the measurement was missing).
self._vf_md_raw: dict[tuple, float] = dict(vf_md_targets or {})
# ── Imputation distribution + pre-sampled imputed table ──────────────
# Strategy: in PAPILA, the few healthy patients with measured MD are
# likely biased toward being slightly worse than the typical healthy
# eye (the test is usually administered when there's some suspicion).
# We correct for this by centering the imputation distribution on the
# TOP QUARTILE mean of the measured-healthy MD values (i.e. the
# healthiest of the measured healthies), while using the std of the
# full measured-healthy sample. Missing values are then drawn from
# N(top_q_mean, std²) once at bundle construction, with a fixed seed,
# so the same patient always gets the same imputed value.
self._impute_mean, self._impute_std = self._compute_impute_params()
self._vf_md_imputed: dict[tuple, float] = self._sample_imputations(seed=42)
# ── Entity-id metadata (for orchestrator logging) ────────────────────────
@property
@@ -526,6 +682,76 @@ class PapilaBundle:
out.append((pid, eye, self.image.get_image_path(pid, eye)))
return out
# ── VF_MD target accessor (for regression / auxiliary heads) ─────────────
def get_vf_md(self, pid: int, eye: str, *, impute: bool = True) -> float:
"""Return VF_MD for one eye.
impute=True → look up the precomputed imputed value (raw if measured,
distribution sample if missing). This is what training
should use — deterministic, fixed-seed, same value every
call across the run.
impute=False → return the raw value (NaN if unmeasured).
"""
key = (int(pid), str(eye))
if impute:
return self._vf_md_imputed.get(key, self._impute_mean)
return self._vf_md_raw.get(key, float("nan"))
@property
def has_vf_md(self) -> bool:
return bool(self._vf_md_raw)
@property
def impute_params(self) -> dict:
"""Inspection: which mean/std were used to draw imputations."""
return {"mean": self._impute_mean, "std": self._impute_std}
# ── Internal helpers for imputation -------------------------------------
def _compute_impute_params(self) -> tuple[float, float]:
"""Mean = top-quartile mean of measured healthy MDs; std = std of all.
Falls back to (0.0, 0.0) if there are too few measured-healthy samples
to fit a sensible distribution.
"""
# Build (pid, eye) → diagnosis from the underlying bundle df.
diag_lookup: dict[tuple, int] = {}
label_col = self._bundle.label_col
pc = self._bundle.patient_col
if label_col in self._bundle.df.columns:
for _, row in self._bundle.df.iterrows():
key = (int(row[pc]), str(row.get("eyeID", "OD")))
diag_lookup[key] = int(row[label_col])
measured_healthy = [
v for key, v in self._vf_md_raw.items()
if v == v and diag_lookup.get(key, -1) == 0 # not NaN and healthy
]
if len(measured_healthy) < 4:
return 0.0, 0.0
arr = np.asarray(measured_healthy, dtype=np.float64)
q3 = float(np.percentile(arr, 75))
top = arr[arr >= q3]
mean = float(top.mean()) if top.size else float(arr.mean())
std = float(arr.std(ddof=1)) if arr.size > 1 else 0.0
return mean, std
def _sample_imputations(self, *, seed: int) -> dict[tuple, float]:
"""One-shot pre-sample: every (pid, eye) gets a fixed value for the run."""
rng = np.random.default_rng(seed)
out: dict[tuple, float] = {}
for key in sorted(self._vf_md_raw.keys()):
v = self._vf_md_raw[key]
if v == v: # measured
out[key] = float(v)
elif self._impute_std > 0:
out[key] = float(rng.normal(self._impute_mean, self._impute_std))
else:
out[key] = self._impute_mean
return out
# ── Backward-compat delegates ────────────────────────────────────────────
@property
@@ -587,7 +813,8 @@ class PapilaBundle:
pid = int(row[pc])
label = int(row[lc])
side = str(row.get("eyeID", "OD"))
entries.append(ShellEntry(entity_id=(pid, side), label=label))
meta = {"vf_md": self.get_vf_md(pid, side)} if self.has_vf_md else {}
entries.append(ShellEntry(entity_id=(pid, side), label=label, meta=meta))
elif level == "patient":
for pid, grp in df.groupby(pc):
@@ -597,7 +824,13 @@ class PapilaBundle:
continue
label_mode = grp[lc].mode()
label = int(label_mode.iloc[0]) if not label_mode.empty else int(grp[lc].iloc[0])
entries.append(ShellEntry(entity_id=(int(pid),), label=label))
meta = {}
if self.has_vf_md:
# Patient-level target: mean of the two eyes' imputed MDs
meta["vf_md"] = 0.5 * (
self.get_vf_md(int(pid), "OD") + self.get_vf_md(int(pid), "OS")
)
entries.append(ShellEntry(entity_id=(int(pid),), label=label, meta=meta))
else:
raise ValueError(f"Unknown shell level: {level!r}. Choose 'eye' or 'patient'.")
@@ -659,6 +892,11 @@ def build_data(args: dict) -> PapilaBundle:
random_seed = int(args.get("random_seed", 42))
use_cache = bool(args.get("in_memory_cache", False))
# Always exclude VF_MD from feature vectorization — it's a target/label
# column (used by regression heads), never an input.
if "VF_MD" not in exclude_cols:
exclude_cols = exclude_cols + ["VF_MD"]
effective_cat = [c for c in cat_cols if c not in exclude_cols]
bundle = DataBundle(
@@ -694,6 +932,20 @@ def build_data(args: dict) -> PapilaBundle:
bundle.frames[i], converter=converter, drop_raw=iop_drop_raw
)
# Capture VF_MD per (pid, eyeID) from the raw frames BEFORE the master df
# is refreshed (which would drop VF_MD via exclude_cols).
vf_md_targets: dict[tuple, float] = {}
for frame in bundle.frames:
if "VF_MD" not in frame.columns:
continue
for _, row in frame.iterrows():
pid = int(row["Patient ID"])
eye = str(row.get("eyeID", "OD"))
v = row["VF_MD"]
vf_md_targets[(pid, eye)] = (
float(v) if not pd.isna(v) else float("nan")
)
bundle._refresh_master_df(exclude_cols=exclude_cols or None)
bundle._infer_or_validate_feature_types(exclude_cols=exclude_cols or None)
bundle._compute_numeric_stats()
@@ -706,4 +958,5 @@ def build_data(args: dict) -> PapilaBundle:
bundle=bundle,
image_dir=image_dir,
image_cache=image_cache,
vf_md_targets=vf_md_targets,
)
v4/classes/stages/fusion.py
+89 -40
View File
@@ -12,7 +12,7 @@ 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 (
class_weights_from_shell, encode_embedding, get_out_dim, resolve_input_dims,
phase_for_epoch,
phase_for_epoch, head_compute_loss, head_to_probs, head_score, head_target_key,
)
@@ -27,18 +27,23 @@ def collect_probs(
device,
num_classes: int,
) -> tuple[np.ndarray, np.ndarray, list, np.ndarray]:
"""Eval pass for one fusion stage; returns (y_true, softmax_probs, entity_ids, embeddings)."""
from v4.classes.dataset import to_label_tensor
"""Eval pass for one fusion stage; returns (y_true, predictions, entity_ids, embeddings).
For classification heads, `predictions` is the softmax over classes (B, C).
For regression heads (or any head with a `to_probs` method), it's whatever
that method returns — typically per-sample scalar predictions.
"""
bridge.eval(); primary_head.eval()
for t in towers.values():
t.eval()
inputs = stage_cfg["inputs"]
is_bilateral = isinstance(inputs, dict)
target_key = head_target_key(primary_head)
y_all, p_all, ids_all, z_all = [], [], [], []
with torch.no_grad():
for batch in loader:
y = batch.get("label")
y = batch.get(target_key, batch.get("label"))
if not torch.is_tensor(y):
continue
@@ -56,13 +61,13 @@ def collect_probs(
z = bridge(embs)
logits = primary_head(z)
y_all.append(to_label_tensor(y, device).cpu().numpy())
p_all.append(F.softmax(logits, dim=1).cpu().numpy())
y_all.append(y.detach().cpu().numpy())
p_all.append(head_to_probs(primary_head, logits))
z_all.append(z.cpu().numpy())
ids_all.extend(batch.get("entity_id", []))
if not y_all:
return (np.zeros(0, dtype=np.int64), np.zeros((0, num_classes), dtype=np.float32),
return (np.zeros(0, dtype=np.float32), np.zeros((0,), dtype=np.float32),
[], np.zeros((0, 0), dtype=np.float32))
return (np.concatenate(y_all), np.concatenate(p_all, axis=0),
ids_all, np.concatenate(z_all, axis=0))
@@ -170,7 +175,10 @@ def run(
warmup_cfg = stage_cfg.get("warmup", {})
wt = 0 if is_bilateral else warmup_cfg.get("tower_epochs", 0)
wf = 0 if is_bilateral else warmup_cfg.get("fused_epochs", 0)
bcd_prob = cfg["training"].get("bcd_prob", 0.5)
bcd_prob = cfg["training"].get("bcd_prob", 0.5)
tower_loss_mode = cfg["training"].get("tower_loss_mode", "bcd")
if tower_loss_mode not in ("bcd", "all_losses"):
raise ValueError(f"tower_loss_mode must be 'bcd' or 'all_losses'; got {tower_loss_mode!r}")
cw = class_weights_from_shell(
s_train, num_classes, device,
@@ -205,6 +213,7 @@ def run(
phase = "fusion"
total_loss = total_correct = total_n = 0
is_class = not hasattr(primary_head, "target_key")
for batch in train_loader:
y = batch.get("label")
@@ -234,10 +243,15 @@ def run(
}
if is_bilateral or phase == "fused_warmup":
logits = head_logits.get(primary_hs_cfg["name"])
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 = [F.cross_entropy(head_logits[hs["name"]], y_t, weight=cw)
for hs in bcd_head_cfgs if hs["name"] in head_logits]
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)
@@ -247,19 +261,42 @@ def run(
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:
logits = head_logits.get(choice(bcd_head_cfgs)["name"])
chosen_hs = choice(bcd_head_cfgs)
else:
logits = head_logits.get(primary_hs_cfg["name"])
chosen_hs = primary_hs_cfg
logits = head_logits.get(chosen_hs["name"])
chosen_head = head_models.get(chosen_hs["name"])
if logits is None:
continue
loss = F.cross_entropy(logits, y_t, weight=cw)
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()
total_correct += int((logits.argmax(1) == y_t).sum())
if logits.dim() >= 2:
total_correct += int((logits.argmax(1) == y_t).sum())
total_loss += loss.item() * len(y_t)
total_n += len(y_t)
@@ -268,49 +305,61 @@ def run(
y_v, p_v, _, _ = collect_probs(bridge, primary_head, stage_cfg, towers,
stage_models, cfg_stages, val_loader, device, num_classes)
_, val_auc, _ = score_arrays(y_v, p_v, num_classes) if y_v.size else (nan, nan, nan)
print(
f" fold{fold+1} [{name}] ep{epoch+1:03d}/{epochs} [{phase:14s}]"
f" loss={tr_loss:.4f} acc={tr_acc:.3f} val_auc={val_auc:.4f}",
flush=True,
)
epoch_scores = head_score(primary_head, y_v, p_v, num_classes)
val_metric = epoch_scores.get("primary", nan)
metric_name = epoch_scores.get("primary_name", "auc")
if is_class:
print(
f" fold{fold+1} [{name}] ep{epoch+1:03d}/{epochs} [{phase:14s}]"
f" loss={tr_loss:.4f} acc={tr_acc:.3f} val_{metric_name}={val_metric:.4f}",
flush=True,
)
else:
print(
f" fold{fold+1} [{name}] ep{epoch+1:03d}/{epochs} [{phase:14s}]"
f" loss={tr_loss:.4f} val_{metric_name}={val_metric:.4f}",
flush=True,
)
# ── Final eval ────────────────────────────────────────────────────────────
y_val, p_val, ids_val, z_val = collect_probs(bridge, primary_head, stage_cfg, towers,
stage_models, cfg_stages, val_loader, device, num_classes)
val_acc, val_auc, val_n = (score_arrays(y_val, p_val, num_classes)
if y_val.size else (nan, nan, nan))
ext = compute_extended_metrics(y_val, p_val, num_classes) if y_val.size else {}
val_scores = head_score(primary_head, y_val, p_val, num_classes)
val_threshold = 0.5
if (cfg["training"].get("tune_binary_threshold")
and num_classes == 2 and y_val.size >= 2):
and num_classes == 2 and y_val.size >= 2
and p_val.ndim == 2 and p_val.shape[1] == 2):
val_threshold = tune_binary_threshold(y_val, p_val[:, 1])
y_te = p_te = ids_te = z_te = None
test_auc = test_acc = test_n = nan
test_scores: dict = {}
if test_loader is not None:
y_te, p_te, ids_te, z_te = collect_probs(bridge, primary_head, stage_cfg, towers,
stage_models, cfg_stages, test_loader, device, num_classes)
test_acc, test_auc, test_n = (score_arrays(y_te, p_te, num_classes)
if y_te.size else (nan, nan, nan))
test_scores = head_score(primary_head, y_te, p_te, num_classes)
updated = dict(stage_models)
updated[name] = bridge
updated.update(head_models)
metrics = {
f"{name}_val_auc": val_auc,
f"{name}_val_acc": val_acc,
f"{name}_val_n": val_n,
f"{name}_val_kappa": ext.get("kappa", nan),
f"{name}_val_mcc": ext.get("mcc", nan),
f"{name}_val_f1": ext.get("macro_f1", nan),
f"{name}_val_threshold": val_threshold,
f"{name}_test_auc": test_auc,
f"{name}_test_acc": test_acc,
f"{name}_test_n": test_n,
}
# Metrics: prefix every score-dict key with stage name + split, plus a fixed
# `_val_primary` / `_test_primary` slot that the summary writer can read
# without knowing the metric's name.
metrics: dict = {f"{name}_val_threshold": val_threshold}
for k, v in val_scores.items():
if k == "primary_name":
continue
if isinstance(v, (int, float)):
metrics[f"{name}_val_{k}"] = float(v)
metrics[f"{name}_val_primary_name"] = val_scores.get("primary_name", "auc")
for k, v in test_scores.items():
if k == "primary_name":
continue
if isinstance(v, (int, float)):
metrics[f"{name}_test_{k}"] = float(v)
metrics[f"{name}_test_primary_name"] = test_scores.get("primary_name",
metrics[f"{name}_val_primary_name"])
pred_data = {
name: {
"val_y": y_val, "val_p": p_val, "val_ids": ids_val, "val_z": z_val,
v4/classes/stages/helpers.py
+74
View File
@@ -3,7 +3,9 @@ from __future__ import annotations
from collections import Counter
import numpy as np
import torch
import torch.nn.functional as F
def class_weights_from_shell(
@@ -87,3 +89,75 @@ def phase_for_epoch(epoch: int, warmup_tower: int, warmup_fused: int) -> str:
if epoch < warmup_tower + warmup_fused:
return "fused_warmup"
return "main"
# ─────────────────────────────────────────────────────────────────────────────
# Head dispatch helpers — opt-in hooks that let new head types (regression,
# ordinal, etc.) plug in without touching the runners. Heads that don't
# implement these methods fall back to the classification defaults.
# ─────────────────────────────────────────────────────────────────────────────
def head_target_key(head) -> str:
"""Which batch field this head consumes as ground truth (default: 'label')."""
return getattr(head, "target_key", "label")
def head_compute_loss(
head, logits: torch.Tensor, batch: dict, y_t: torch.Tensor,
*, class_weights: torch.Tensor | None = None,
) -> torch.Tensor:
"""Compute one head's training loss.
If the head provides `compute_loss(logits, batch)`, that wins — the head
is responsible for reading its own target from batch and applying whatever
loss function it wants. Otherwise we fall back to weighted CE against y_t
(the standard classification path).
"""
if hasattr(head, "compute_loss"):
return head.compute_loss(logits, batch)
return F.cross_entropy(logits, y_t, weight=class_weights)
def head_to_probs(head, logits: torch.Tensor) -> np.ndarray:
"""Convert raw head outputs to a per-sample numpy array.
For classification heads, this is the softmax over class logits.
For regression heads, this is just the raw predicted scalar(s).
Heads override `to_probs(logits)` to define their own conversion.
"""
if hasattr(head, "to_probs"):
return head.to_probs(logits)
return F.softmax(logits, dim=1).cpu().numpy()
def head_score(head, y_true: np.ndarray, predictions: np.ndarray,
num_classes: int) -> dict:
"""Compute evaluation metrics for one head.
Returns a dict with at minimum:
- 'primary' : the headline metric value (float)
- 'primary_name' : how to label it (e.g. 'auc', 'mse')
- 'n' : sample count
Classification heads fall through to the existing score_arrays-based
metric set. Regression heads override `score(y_true, predictions)` and
return their own dict (which may include the standard fields plus extras
like 'mae', 'r2', 'spearman').
"""
if hasattr(head, "score"):
return head.score(y_true, predictions)
# Default: classification scoring
from v4.classes.metrics import score_arrays, compute_extended_metrics
if not y_true.size:
return {"primary": float("nan"), "primary_name": "auc",
"auc": float("nan"), "acc": float("nan"), "n": 0}
acc, auc, n = score_arrays(y_true, predictions, num_classes)
ext = compute_extended_metrics(y_true, predictions, num_classes)
return {
"primary": float(auc),
"primary_name": "auc",
"auc": float(auc),
"acc": float(acc),
"n": int(n),
**{k: float(v) for k, v in ext.items() if isinstance(v, (int, float))},
}
v4/classes/stages/parallel.py
+94 -41
View File
@@ -16,6 +16,7 @@ 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 (
class_weights_from_shell, get_out_dim, resolve_input_dims, phase_for_epoch,
head_compute_loss, head_to_probs, head_score, head_target_key,
)
from v4.classes.stages.fusion import collect_probs
@@ -120,18 +121,31 @@ def _parallel_warm(
continue
y_t = to_label_tensor(y, device)
logits = ctx["probe"](towers[tower_name](x.to(device)))
loss = F.cross_entropy(logits, y_t, weight=ctx["class_weights"])
loss = head_compute_loss(ctx["probe"], logits, batch, y_t,
class_weights=ctx["class_weights"])
ctx["opt"].zero_grad(); loss.backward(); ctx["opt"].step()
total_loss += loss.item() * len(y_t)
total_correct += int((logits.argmax(1) == y_t).sum())
if logits.dim() >= 2:
total_correct += int((logits.argmax(1) == y_t).sum())
is_class = True
else:
is_class = False
total_n += len(y_t)
if total_n:
print(
f" fold{fold+1} [warm/{tower_name}]"
f" ep{epoch+1:03d}/{ctx['n_epochs']}"
f" loss={total_loss/total_n:.4f} acc={total_correct/total_n:.3f}",
flush=True,
)
if is_class:
print(
f" fold{fold+1} [warm/{tower_name}]"
f" ep{epoch+1:03d}/{ctx['n_epochs']}"
f" loss={total_loss/total_n:.4f} acc={total_correct/total_n:.3f}",
flush=True,
)
else:
print(
f" fold{fold+1} [warm/{tower_name}]"
f" ep{epoch+1:03d}/{ctx['n_epochs']}"
f" loss={total_loss/total_n:.4f}",
flush=True,
)
for t in towers.values():
for p in t.parameters():
@@ -153,7 +167,10 @@ def _parallel_fusion(
):
nan = float("nan")
bs = cfg["training"]["batch_size"]
bcd_prob = cfg["training"].get("bcd_prob", 0.5)
bcd_prob = cfg["training"].get("bcd_prob", 0.5)
tower_loss_mode = cfg["training"].get("tower_loss_mode", "bcd")
if tower_loss_mode not in ("bcd", "all_losses"):
raise ValueError(f"tower_loss_mode must be 'bcd' or 'all_losses'; got {tower_loss_mode!r}")
# Freeze all prior stage models once, before building any bridges.
for m in stage_models.values():
@@ -302,10 +319,14 @@ def _parallel_fusion(
if phase == "fused_warmup":
logits = head_logits.get(ctx["primary_hs_cfg"]["name"])
chosen_head = ctx["head_models"].get(ctx["primary_hs_cfg"]["name"])
elif phase == "tower_warmup" and ctx["bcd_head_cfgs"]:
losses = [F.cross_entropy(head_logits[hs["name"]], y_t,
weight=ctx["class_weights"])
for hs in ctx["bcd_head_cfgs"] if hs["name"] in head_logits]
losses = [
head_compute_loss(ctx["head_models"][hs["name"]],
head_logits[hs["name"]], batch, y_t,
class_weights=ctx["class_weights"])
for hs in ctx["bcd_head_cfgs"] if hs["name"] in head_logits
]
if not losses:
continue
loss = sum(losses) / len(losses)
@@ -315,19 +336,40 @@ def _parallel_fusion(
total_loss += loss.item() * len(y_t)
total_n += len(y_t)
continue
elif tower_loss_mode == "all_losses" and ctx["bcd_head_cfgs"]:
all_head_names = ([ctx["primary_hs_cfg"]["name"]]
+ [hs["name"] for hs in ctx["bcd_head_cfgs"]])
losses = [
head_compute_loss(ctx["head_models"][n], head_logits[n],
batch, y_t, class_weights=ctx["class_weights"])
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)
ctx["opt"].zero_grad(); loss.backward(); ctx["opt"].step()
total_loss += loss.item() * len(y_t)
total_n += len(y_t)
continue
else:
if ctx["bcd_head_cfgs"] and _random() < bcd_prob:
logits = head_logits.get(choice(ctx["bcd_head_cfgs"])["name"])
chosen_hs = choice(ctx["bcd_head_cfgs"])
else:
logits = head_logits.get(ctx["primary_hs_cfg"]["name"])
chosen_hs = ctx["primary_hs_cfg"]
logits = head_logits.get(chosen_hs["name"])
chosen_head = ctx["head_models"].get(chosen_hs["name"])
if logits is None:
continue
loss = F.cross_entropy(logits, y_t, weight=ctx["class_weights"])
loss = head_compute_loss(chosen_head, logits, batch, y_t,
class_weights=ctx["class_weights"])
if hasattr(bridge, "modify_loss"):
loss = bridge.modify_loss(loss)
ctx["opt"].zero_grad(); loss.backward(); ctx["opt"].step()
total_correct += int((logits.argmax(1) == y_t).sum())
if logits.dim() >= 2:
total_correct += int((logits.argmax(1) == y_t).sum())
total_loss += loss.item() * len(y_t)
total_n += len(y_t)
@@ -337,12 +379,22 @@ def _parallel_fusion(
y_v, p_v, _, _ = collect_probs(bridge, ctx["primary_head"], ctx["sc"], towers,
stage_models, cfg_stages, ctx["val_loader"],
device, num_classes)
_, val_auc, _ = score_arrays(y_v, p_v, num_classes) if y_v.size else (nan, nan, nan)
print(
f" fold{fold+1} [{name}] ep{epoch+1:03d}/{ctx['epochs']} [{phase:14s}]"
f" loss={tr_loss:.4f} acc={tr_acc:.3f} val_auc={val_auc:.4f}",
flush=True,
)
epoch_scores = head_score(ctx["primary_head"], y_v, p_v, num_classes)
val_metric = epoch_scores.get("primary", nan)
metric_name = epoch_scores.get("primary_name", "auc")
is_class = not hasattr(ctx["primary_head"], "target_key")
if is_class:
print(
f" fold{fold+1} [{name}] ep{epoch+1:03d}/{ctx['epochs']} [{phase:14s}]"
f" loss={tr_loss:.4f} acc={tr_acc:.3f} val_{metric_name}={val_metric:.4f}",
flush=True,
)
else:
print(
f" fold{fold+1} [{name}] ep{epoch+1:03d}/{ctx['epochs']} [{phase:14s}]"
f" loss={tr_loss:.4f} val_{metric_name}={val_metric:.4f}",
flush=True,
)
# ── Final eval + collect results ──────────────────────────────────────────
updated = dict(stage_models)
@@ -360,36 +412,37 @@ def _parallel_fusion(
y_val, p_val, ids_val, z_val = collect_probs(bridge, primary_head, ctx["sc"], towers,
stage_models, cfg_stages, ctx["val_loader"],
device, num_classes)
val_acc, val_auc, val_n = (score_arrays(y_val, p_val, num_classes)
if y_val.size else (nan, nan, nan))
ext = compute_extended_metrics(y_val, p_val, num_classes) if y_val.size else {}
val_scores = head_score(primary_head, y_val, p_val, num_classes)
val_threshold = 0.5
if (cfg["training"].get("tune_binary_threshold")
and num_classes == 2 and y_val.size >= 2):
and num_classes == 2 and y_val.size >= 2
and p_val.ndim == 2 and p_val.shape[1] == 2):
val_threshold = tune_binary_threshold(y_val, p_val[:, 1])
y_te = p_te = ids_te = z_te = None
test_auc = test_acc = test_n = nan
test_scores: dict = {}
if ctx["test_loader"] is not None:
y_te, p_te, ids_te, z_te = collect_probs(bridge, primary_head, ctx["sc"], towers,
stage_models, cfg_stages, ctx["test_loader"],
device, num_classes)
test_acc, test_auc, test_n = (score_arrays(y_te, p_te, num_classes)
if y_te.size else (nan, nan, nan))
test_scores = head_score(primary_head, y_te, p_te, num_classes)
all_metrics.update({
f"{name}_val_auc": val_auc,
f"{name}_val_acc": val_acc,
f"{name}_val_n": val_n,
f"{name}_val_kappa": ext.get("kappa", nan),
f"{name}_val_mcc": ext.get("mcc", nan),
f"{name}_val_f1": ext.get("macro_f1", nan),
f"{name}_val_threshold": val_threshold,
f"{name}_test_auc": test_auc,
f"{name}_test_acc": test_acc,
f"{name}_test_n": test_n,
})
per_stage_metrics: dict = {f"{name}_val_threshold": val_threshold}
for k, v in val_scores.items():
if k == "primary_name":
continue
if isinstance(v, (int, float)):
per_stage_metrics[f"{name}_val_{k}"] = float(v)
per_stage_metrics[f"{name}_val_primary_name"] = val_scores.get("primary_name", "auc")
for k, v in test_scores.items():
if k == "primary_name":
continue
if isinstance(v, (int, float)):
per_stage_metrics[f"{name}_test_{k}"] = float(v)
per_stage_metrics[f"{name}_test_primary_name"] = test_scores.get(
"primary_name", per_stage_metrics[f"{name}_val_primary_name"])
all_metrics.update(per_stage_metrics)
all_preds[name] = {
"val_y": y_val, "val_p": p_val, "val_ids": ids_val, "val_z": z_val,
"test_y": y_te, "test_p": p_te, "test_ids": ids_te, "test_z": z_te,
v4/classes/stages/warm.py
+19 -8
View File
@@ -7,7 +7,7 @@ import torch
import torch.nn.functional as F
from v4.classes.dataset import to_label_tensor
from v4.classes.stages.helpers import class_weights_from_shell
from v4.classes.stages.helpers import class_weights_from_shell, head_compute_loss
def run(
@@ -80,6 +80,7 @@ def run(
towers[tower_name].train()
probe.train()
is_classification = True
for epoch in range(n_epochs):
total_loss = total_correct = total_n = 0
for batch in loader:
@@ -89,16 +90,26 @@ def run(
continue
y_t = to_label_tensor(y, device)
logits = probe(towers[tower_name](x.to(device)))
loss = F.cross_entropy(logits, y_t, weight=cw)
loss = head_compute_loss(probe, logits, batch, y_t, class_weights=cw)
opt.zero_grad(); loss.backward(); opt.step()
total_loss += loss.item() * len(y_t)
total_correct += int((logits.argmax(1) == y_t).sum())
if logits.dim() >= 2:
total_correct += int((logits.argmax(1) == y_t).sum())
else:
is_classification = False
total_n += len(y_t)
print(
f" fold{fold+1} [warm/{tower_name}] ep{epoch+1:03d}/{n_epochs}"
f" loss={total_loss/total_n:.4f} acc={total_correct/total_n:.3f}",
flush=True,
)
if is_classification:
print(
f" fold{fold+1} [warm/{tower_name}] ep{epoch+1:03d}/{n_epochs}"
f" loss={total_loss/total_n:.4f} acc={total_correct/total_n:.3f}",
flush=True,
)
else:
print(
f" fold{fold+1} [warm/{tower_name}] ep{epoch+1:03d}/{n_epochs}"
f" loss={total_loss/total_n:.4f}",
flush=True,
)
for t in towers.values():
for p in t.parameters():
v4/classes/v4_hypertower.py
+61 -18
View File
@@ -21,6 +21,7 @@ from pathlib import Path
from typing import Any
import numpy as np
import torch
from torch.utils.data import DataLoader
REPO_ROOT = Path(__file__).resolve().parents[2]
@@ -148,7 +149,13 @@ def _balanced_sampler(shell: LoaderShell):
# Fold runner
# ---------------------------------------------------------------------------
def run_fold(fold: int, splits, cfg: dict, data, num_classes: int, device) -> dict:
def run_fold(fold: int, splits, cfg: dict, data, num_classes: int, device):
"""Train + evaluate one fold.
Returns (fold_result, fold_preds, towers, stage_models). The latter two are
handy for opt-in artefact saving (e.g. checkpoint dumps for explainability
runs) without forcing the orchestrator to know about every saved tensor.
"""
seed_everything(cfg["seed"] + fold * 100)
split = splits[fold]
label_filter = cfg.get("label_filter", None)
@@ -210,7 +217,7 @@ def run_fold(fold: int, splits, cfg: dict, data, num_classes: int, device) -> di
# head stages are handled inside fusion.run
return fold_result, fold_preds
return fold_result, fold_preds, towers, stage_models
# ---------------------------------------------------------------------------
@@ -264,6 +271,7 @@ def main():
eval_stage = cfg.get("eval_stage", "hb")
save_predictions = cfg.get("save_predictions", False)
save_features = cfg.get("save_features", False)
save_checkpoints = cfg.get("save_checkpoints", False)
fold_results = []
eval_stage_preds = [] # list[dict] — one per fold, only for eval_stage
all_phase_preds: dict[str, list[dict]] = {} # phase → list[dict] across folds
@@ -274,8 +282,21 @@ def main():
n_train = split.train[group_col].nunique() if group_col else len(split.train)
print(f"\n── fold {fold+1}/{cfg.get('folds', 5)} train_groups={n_train} ──",
flush=True)
result, fold_preds = run_fold(fold, splits, cfg, data, num_classes, device)
result, fold_preds, towers, stage_models = run_fold(
fold, splits, cfg, data, num_classes, device,
)
fold_results.append(result)
if save_checkpoints:
ckpt_dir = out_dir / "checkpoints" / f"fold{fold}"
ckpt_dir.mkdir(parents=True, exist_ok=True)
for name, mod in towers.items():
torch.save(mod.state_dict(), ckpt_dir / f"tower_{name}.pt")
for name, mod in stage_models.items():
# Skip non-Module entries (defensive); only nn.Modules have state_dict
if hasattr(mod, "state_dict"):
torch.save(mod.state_dict(), ckpt_dir / f"stage_{name}.pt")
print(f" Checkpoints saved: {ckpt_dir}", flush=True)
if save_predictions and eval_stage in fold_preds:
eval_stage_preds.append(fold_preds[eval_stage])
if save_features:
@@ -283,26 +304,46 @@ def main():
if pdata.get("val_z") is None:
continue
all_phase_preds.setdefault(ph, []).append(pdata)
# Look up the eval stage's primary metric name (set by the stage runner).
primary_name = result.get(f"{eval_stage}_val_primary_name", "auc")
val_primary = result.get(f"{eval_stage}_val_{primary_name}", float("nan"))
test_primary = result.get(f"{eval_stage}_test_{primary_name}", float("nan"))
print(
f" fold{fold+1} DONE"
f" val_auc={result.get(f'{eval_stage}_val_auc', float('nan')):.4f}"
f" test_auc={result.get(f'{eval_stage}_test_auc', float('nan')):.4f}",
f" val_{primary_name}={val_primary:.4f}"
f" test_{primary_name}={test_primary:.4f}",
flush=True,
)
if fold_results:
val_aucs = [r.get(f"{eval_stage}_val_auc", float("nan")) for r in fold_results]
test_aucs = [r.get(f"{eval_stage}_test_auc", float("nan")) for r in fold_results]
val_aucs = [v for v in val_aucs if not np.isnan(v)]
test_aucs = [v for v in test_aucs if not np.isnan(v)]
# Resolve primary metric name from first valid fold result.
primary_name = next(
(r.get(f"{eval_stage}_val_primary_name", "auc") for r in fold_results
if r.get(f"{eval_stage}_val_primary_name") is not None),
"auc",
)
val_primaries = [r.get(f"{eval_stage}_val_{primary_name}", float("nan"))
for r in fold_results]
test_primaries = [r.get(f"{eval_stage}_test_{primary_name}", float("nan"))
for r in fold_results]
val_primaries = [v for v in val_primaries if not np.isnan(v)]
test_primaries = [v for v in test_primaries if not np.isnan(v)]
summary = {
"run_name": cfg["run_name"],
"eval_stage": eval_stage,
"config": cfg,
"mean_val_auc": float(np.mean(val_aucs)) if val_aucs else float("nan"),
"std_val_auc": float(np.std(val_aucs)) if val_aucs else float("nan"),
"mean_test_auc": float(np.mean(test_aucs)) if test_aucs else float("nan"),
"std_test_auc": float(np.std(test_aucs)) if test_aucs else float("nan"),
"run_name": cfg["run_name"],
"eval_stage": eval_stage,
"primary_metric": primary_name,
"config": cfg,
# Canonical primary-metric stats
f"mean_val_{primary_name}": float(np.mean(val_primaries)) if val_primaries else float("nan"),
f"std_val_{primary_name}": float(np.std(val_primaries)) if val_primaries else float("nan"),
f"mean_test_{primary_name}": float(np.mean(test_primaries)) if test_primaries else float("nan"),
f"std_test_{primary_name}": float(np.std(test_primaries)) if test_primaries else float("nan"),
# Backward-compat aliases so existing analysis tooling (summarize_run.py,
# compare_grid.py) still reads correctly for classification runs.
"mean_val_auc": float(np.mean(val_primaries)) if val_primaries else float("nan"),
"std_val_auc": float(np.std(val_primaries)) if val_primaries else float("nan"),
"mean_test_auc": float(np.mean(test_primaries)) if test_primaries else float("nan"),
"std_test_auc": float(np.std(test_primaries)) if test_primaries else float("nan"),
"elapsed_s": round(time.time() - t0, 1),
"fold_results": fold_results,
}
@@ -310,8 +351,10 @@ def main():
summary_path = out_dir / "summary.json"
summary_path.write_text(json.dumps(summary, indent=2))
print(f"\n{'='*60}", flush=True)
print(f"Val AUC: {summary['mean_val_auc']:.4f} ± {summary['std_val_auc']:.4f}", flush=True)
print(f"Test AUC: {summary['mean_test_auc']:.4f} ± {summary['std_test_auc']:.4f}", flush=True)
print(f"Val {primary_name}: {summary[f'mean_val_{primary_name}']:.4f} ± "
f"{summary[f'std_val_{primary_name}']:.4f}", flush=True)
print(f"Test {primary_name}: {summary[f'mean_test_{primary_name}']:.4f} ± "
f"{summary[f'std_test_{primary_name}']:.4f}", flush=True)
print(f"Saved: {summary_path}", flush=True)
if save_predictions and eval_stage_preds:
v4/classes/vehicles-deprecated/__init__.py
View File
v4/classes/vehicles-deprecated/htbase.py
-40
View File
@@ -1,40 +0,0 @@
"""htbase — HTBase: abstract base for all v4 vehicle classes."""
from __future__ import annotations
from abc import ABC, abstractmethod
import torch
import torch.nn as nn
class HTBase(nn.Module, ABC):
"""Shared interface for all HyperTower vehicles.
Subclasses must implement ``encode`` and ``forward``.
``transform`` walks ``self.towers`` (if present) and returns the transform
from the first tower that exposes one — used by data loaders.
``forward`` contract: returns ``(logits, aux_dict)`` where
``aux_dict`` maps a name or index to per-component logits.
HTMono returns an empty dict to keep the signature uniform.
"""
@abstractmethod
def encode(self, inputs) -> torch.Tensor:
"""Return the pre-classifier embedding."""
@abstractmethod
def forward(self, inputs) -> tuple[torch.Tensor, dict]:
"""Return (logits, aux_dict)."""
@property
def transform(self):
towers = getattr(self, "towers", None) or {}
for t in (towers.values() if hasattr(towers, "values") else []):
if hasattr(t, "transform"):
return t.transform
encoder = getattr(self, "encoder", None)
if encoder is not None:
return getattr(encoder, "transform", None)
return None
v4/classes/vehicles-deprecated/htfusion.py
-61
View File
@@ -1,61 +0,0 @@
"""htfusion — HTFusion: N named towers fused through a FusionBridge."""
from __future__ import annotations
import torch
import torch.nn as nn
from v4.classes.bridges.fusion_bridge import FusionBridge
from v4.classes.vehicles.htbase import HTBase
class HTFusion(HTBase):
"""General N-tower fusion vehicle.
Each named encoder is registered as a submodule; the FusionBridge
projects and Hadamard-fuses their embeddings.
Parameters
----------
towers : ordered dict ``{name: encoder}``. Each encoder must
expose ``.out_dim``.
num_classes : output classes
fusion_dim : bridge projection dimensionality
dropout : bridge dropout
use_se : SE gate on the fused vector
Forward contract
----------------
``forward(embeddings)`` takes a ``dict[str, Tensor]`` of pre-computed
per-tower embeddings and returns ``(logits_fused, aux_dict)`` where
``aux_dict`` maps each tower name to its auxiliary head logits.
"""
def __init__(
self,
towers: dict[str, nn.Module],
num_classes: int,
fusion_dim: int = 256,
dropout: float = 0.5,
use_se: bool = False,
):
super().__init__()
self.towers = nn.ModuleDict(towers)
self.bridge = FusionBridge(
tower_dims=[t.out_dim for t in self.towers.values()],
num_classes=num_classes,
fusion_dim=fusion_dim,
dropout=dropout,
use_se=use_se,
)
def encode(self, embeddings: dict[str, torch.Tensor]) -> torch.Tensor:
"""Return z_fused (pre-classifier) from a dict of per-tower embeddings."""
return self.bridge.encode([embeddings[name] for name in self.towers])
def forward(
self,
embeddings: dict[str, torch.Tensor],
) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
ordered = [embeddings[name] for name in self.towers]
logits, aux = self.bridge.fuse(ordered)
return logits, {name: aux[i] for i, name in enumerate(self.towers)}
v4/classes/vehicles-deprecated/htlateral.py
-67
View File
@@ -1,67 +0,0 @@
"""htlateral — HTLateral: shared encoder over N same-type inputs."""
from __future__ import annotations
import torch
import torch.nn as nn
from v4.classes.vehicles.htbase import HTBase
class HTLateral(HTBase):
"""N same-type inputs through a shared encoder, jointly compressed, then classified.
All inputs share the same encoder weights (one forward pass per input).
The joint MLP compresses the concatenated embeddings before classification.
Aux heads provide per-input logits before the joint MLP — useful for
BCD-style training.
Parameters
----------
encoder : shared encoder module with ``.out_dim``
input_names : ordered slot names (e.g. ``["od", "os"]``)
num_classes : output classes
fusion_dim : joint MLP hidden dim
dropout : dropout in MLP and classifier
"""
def __init__(
self,
encoder: nn.Module,
input_names: list[str],
num_classes: int,
fusion_dim: int = 256,
dropout: float = 0.5,
):
super().__init__()
self.encoder = encoder
self.input_names = list(input_names)
n = len(input_names)
in_dim: int = encoder.out_dim # type: ignore[assignment]
self.joint = nn.Sequential(
nn.Linear(n * in_dim, fusion_dim), nn.LayerNorm(fusion_dim),
nn.ReLU(), nn.Dropout(dropout), nn.Linear(fusion_dim, in_dim),
)
self.aux_heads = nn.ModuleList([
nn.Linear(in_dim, num_classes) for _ in range(n)
])
self.head = nn.Sequential(
nn.ReLU(), nn.Dropout(dropout), nn.Linear(in_dim, num_classes),
)
def encode(self, inputs: dict[str, torch.Tensor]) -> torch.Tensor:
"""Return joint embedding (post-MLP, pre-classifier)."""
zs = [self.encoder(inputs[name]) for name in self.input_names]
return self.joint(torch.cat(zs, dim=1))
def forward(
self,
inputs: dict[str, torch.Tensor],
) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
zs = [self.encoder(inputs[name]) for name in self.input_names]
z_joint = self.joint(torch.cat(zs, dim=1))
logits = self.head(z_joint)
aux = {name: head(z)
for name, head, z in zip(self.input_names, self.aux_heads, zs)}
return logits, aux
v4/classes/vehicles-deprecated/htmono.py
-44
View File
@@ -1,44 +0,0 @@
"""htmono — HTMono: single tower + ClassificationHead, no bridge."""
from __future__ import annotations
import torch
import torch.nn as nn
from v4.classes.heads.classifier import ClassificationHead
from v4.classes.vehicles.htbase import HTBase
class HTMono(HTBase):
"""Single-tower vehicle: tower embedding fed directly into a ClassificationHead.
No bridge or projection — the tower's output goes straight to
ReLU → Dropout → Linear. Returns ``(logits, {})`` from ``forward``
to match the HTFusion / HTLateral interface.
Parameters
----------
tower : encoder module with ``.out_dim``
num_classes : output classes
dropout : dropout before the output linear layer
"""
def __init__(
self,
tower: nn.Module,
num_classes: int,
dropout: float = 0.5,
):
super().__init__()
self.tower = tower
self.head = ClassificationHead(tower.out_dim, num_classes, dropout) # type: ignore[arg-type]
def encode(self, inputs) -> torch.Tensor:
"""Return tower embedding (pre-classifier)."""
if isinstance(inputs, dict):
# single-entry dict from HTDataset eye-level pass
(z,) = inputs.values()
return self.tower(z) if torch.is_tensor(z) else self.tower(*z.values())
return self.tower(inputs)
def forward(self, inputs) -> tuple[torch.Tensor, dict]:
return self.head(self.encode(inputs)), {}