Add new regression and ensemble experiment configurations for V2-M and OrthoBridge
- 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
+89
-40
@@ -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,
|
||||
|
||||
@@ -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))},
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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():
|
||||
|
||||
Reference in New Issue
Block a user