rpotter6298
1697 lines
76 KiB
Python
Executable File
1697 lines
76 KiB
Python
Executable File
from __future__ import annotations
|
|
import os
|
|
import csv
|
|
import math
|
|
import argparse
|
|
import json
|
|
from pathlib import Path
|
|
import shutil
|
|
import logging
|
|
import pandas as pd
|
|
from PIL import Image, ImageDraw
|
|
import torch
|
|
from torch import nn
|
|
from torch.utils.data import DataLoader
|
|
from torch.utils.data.sampler import WeightedRandomSampler
|
|
from classes import ClinicalData, ClinicalDataset, ImageTower, MDTower, Bridge, VoteBridge, EarlyStopper
|
|
from classes.unet_segmenter import UNetSegmenter
|
|
from classes.geometry_features import (
|
|
FEATURE_DIM,
|
|
compute_geometry_features,
|
|
disc_cup_from_mask_image,
|
|
)
|
|
from classes.refuge_classification import _geometry_from_mask
|
|
from torchvision import transforms
|
|
|
|
# from clinical_data import ClinicalData
|
|
# from dataset import ClinicalDataset
|
|
# from image_tower import ImageTower
|
|
# from md_tower import MDTower
|
|
# from bridge import Bridge, VoteBridge
|
|
from random import random
|
|
from sklearn.metrics import roc_auc_score
|
|
import torch.nn.functional as F
|
|
import numpy as np
|
|
from sklearn.metrics import roc_curve, auc
|
|
from sklearn.preprocessing import label_binarize
|
|
from typing import Optional, Tuple
|
|
|
|
|
|
LOG_FIELDS = [
|
|
"epoch",
|
|
"eval_loss",
|
|
"acc_fused", "acc_img", "acc_md",
|
|
"auc_fused", "auc_img", "auc_md",
|
|
"top2_fused", "top2_img", "top2_md",
|
|
"margin_fused", "margin_img", "margin_md",
|
|
"agree_fused_img", "agree_fused_md",
|
|
"pct_fused", "pct_img", "pct_md",
|
|
"phase",
|
|
"holdout_loss",
|
|
"holdout_acc_fused",
|
|
"holdout_acc_img",
|
|
"holdout_acc_md",
|
|
"holdout_auc_fused",
|
|
"holdout_auc_img",
|
|
"holdout_auc_md",
|
|
]
|
|
|
|
|
|
def focal_loss(
|
|
logits: torch.Tensor,
|
|
targets: torch.Tensor,
|
|
gamma: float = 0.0,
|
|
weight: Optional[torch.Tensor] = None,
|
|
reduction: str = "mean",
|
|
) -> torch.Tensor:
|
|
"""
|
|
Standard focal loss wrapper. When gamma=0 it reduces to cross entropy.
|
|
weight should be per-class weights (same semantics as CrossEntropyLoss).
|
|
"""
|
|
if gamma <= 0:
|
|
return F.cross_entropy(logits, targets, weight=weight, reduction=reduction)
|
|
|
|
log_probs = F.log_softmax(logits, dim=1)
|
|
probs = log_probs.exp()
|
|
|
|
targets = targets.long().view(-1, 1)
|
|
logpt = log_probs.gather(1, targets)
|
|
pt = probs.gather(1, targets)
|
|
|
|
focal_factor = (1.0 - pt).clamp_min(0.0) ** gamma
|
|
loss = -focal_factor * logpt
|
|
|
|
if weight is not None:
|
|
class_weight = weight.gather(0, targets.view(-1))
|
|
loss = loss * class_weight.view(-1, 1)
|
|
|
|
loss = loss.view(-1)
|
|
if reduction == "sum":
|
|
return loss.sum()
|
|
if reduction == "mean":
|
|
return loss.mean()
|
|
return loss
|
|
|
|
|
|
class UNetImageCropper:
|
|
def __init__(
|
|
self,
|
|
manifest_path: Path,
|
|
weights_path: Path,
|
|
normalize: str = "per_image",
|
|
threshold: float = 0.5,
|
|
tta: bool = False,
|
|
scale: float = 2.5,
|
|
target_size: int = 224,
|
|
cache_dir: Optional[Path] = None,
|
|
) -> None:
|
|
self.segmenter = UNetSegmenter(
|
|
manifest_path=manifest_path,
|
|
normalize=normalize,
|
|
)
|
|
state = torch.load(weights_path, map_location=self.segmenter.device)
|
|
state_dict = state.get("model", state)
|
|
self.segmenter.model.load_state_dict(state_dict)
|
|
self.segmenter.model.to(self.segmenter.device)
|
|
self.segmenter.model.eval()
|
|
|
|
self.threshold = threshold
|
|
self.tta = tta
|
|
self.scale = scale
|
|
self.target_size = target_size
|
|
self.cache_dir = Path(cache_dir) if cache_dir is not None else None
|
|
if self.cache_dir is not None:
|
|
self.cache_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
self.to_tensor = transforms.ToTensor()
|
|
|
|
def _cache_path(self, image_path: Path) -> Optional[Path]:
|
|
if self.cache_dir is None:
|
|
return None
|
|
stem = image_path.stem
|
|
return self.cache_dir / f"{stem}_s{int(self.scale * 100)}.npz"
|
|
|
|
def _infer_masks(self, image: Image.Image) -> Optional[Tuple[np.ndarray, np.ndarray]]:
|
|
resized = self.segmenter.preprocess_image(image)
|
|
tensor = self.to_tensor(resized).unsqueeze(0).to(self.segmenter.device)
|
|
|
|
with torch.no_grad():
|
|
logits = self.segmenter.model(tensor)
|
|
if self.tta:
|
|
t_h = torch.flip(tensor, dims=[3])
|
|
log_h = self.segmenter.model(t_h)
|
|
log_h = torch.flip(log_h, dims=[3])
|
|
t_v = torch.flip(tensor, dims=[2])
|
|
log_v = self.segmenter.model(t_v)
|
|
log_v = torch.flip(log_v, dims=[2])
|
|
logits = (logits + log_h + log_v) / 3.0
|
|
probs = torch.sigmoid(logits)[0].cpu().numpy()
|
|
|
|
disc_pred = (probs[0] > self.threshold).astype(np.uint8) * 255
|
|
cup_pred = (probs[1] > self.threshold).astype(np.uint8) * 255
|
|
disc_img = Image.fromarray(disc_pred, mode="L").resize(image.size, Image.NEAREST)
|
|
disc_mask = np.array(disc_img, dtype=np.uint8)
|
|
cup_img = Image.fromarray(cup_pred, mode="L").resize(image.size, Image.NEAREST)
|
|
cup_mask = (np.array(cup_img, dtype=np.uint8) > 0).astype(np.uint8)
|
|
cup_mask = (cup_mask > 0) & (disc_mask > 0)
|
|
cup_mask = cup_mask.astype(np.uint8)
|
|
disc_mask = (disc_mask > 0).astype(np.uint8)
|
|
return disc_mask, cup_mask
|
|
|
|
def _compute_crop_info(self, image: Image.Image, image_path: Path) -> Optional[dict]:
|
|
image_path = Path(image_path).resolve()
|
|
cache_path = self._cache_path(image_path)
|
|
cached_bounds = None
|
|
if cache_path is not None and cache_path.exists():
|
|
data = np.load(cache_path, allow_pickle=False)
|
|
try:
|
|
cached_bounds = {
|
|
"left": float(data["left"]),
|
|
"upper": float(data["upper"]),
|
|
"right": float(data["right"]),
|
|
"lower": float(data["lower"]),
|
|
}
|
|
if "features" in data.files:
|
|
cached_bounds["features"] = data["features"].astype(np.float32)
|
|
return cached_bounds
|
|
except KeyError:
|
|
cached_bounds = None
|
|
|
|
masks = self._infer_masks(image)
|
|
if masks is None:
|
|
return cached_bounds
|
|
disc_mask, cup_mask = masks
|
|
try:
|
|
geom = _geometry_from_mask(disc_mask, self.scale)
|
|
except Exception:
|
|
return cached_bounds
|
|
cx = geom["centre_x"]
|
|
cy = geom["centre_y"]
|
|
r = geom["crop_radius"]
|
|
left = max(0.0, cx - r)
|
|
upper = max(0.0, cy - r)
|
|
right = min(float(image.width), cx + r)
|
|
lower = min(float(image.height), cy + r)
|
|
features = compute_geometry_features(disc_mask, cup_mask)
|
|
|
|
info = {
|
|
"left": left,
|
|
"upper": upper,
|
|
"right": right,
|
|
"lower": lower,
|
|
"features": features,
|
|
}
|
|
if cache_path is not None:
|
|
np.savez(
|
|
cache_path,
|
|
left=left,
|
|
upper=upper,
|
|
right=right,
|
|
lower=lower,
|
|
width=float(image.width),
|
|
height=float(image.height),
|
|
scale=self.scale,
|
|
target_size=self.target_size,
|
|
features=features,
|
|
)
|
|
return info
|
|
|
|
def __call__(self, image: Image.Image, image_path: Path) -> Image.Image:
|
|
info = self._compute_crop_info(image, image_path)
|
|
if info is None:
|
|
return image
|
|
left = info["left"]
|
|
upper = info["upper"]
|
|
right = info["right"]
|
|
lower = info["lower"]
|
|
if right <= left or lower <= upper:
|
|
return image
|
|
crop = image.crop((left, upper, right, lower))
|
|
return crop.resize((self.target_size, self.target_size), Image.BILINEAR)
|
|
|
|
def geometry_features(self, image: Image.Image, image_path: Path) -> Optional[np.ndarray]:
|
|
info = self._compute_crop_info(image, image_path)
|
|
if info is None:
|
|
return None
|
|
features = info.get("features")
|
|
if features is None:
|
|
return None
|
|
return np.asarray(features, dtype=np.float32)
|
|
|
|
|
|
class ManifestImageCropper:
|
|
def __init__(
|
|
self,
|
|
manifest_path: Path,
|
|
scale: float = 2.5,
|
|
target_size: int = 224,
|
|
cache_dir: Optional[Path] = None,
|
|
) -> None:
|
|
self.scale = scale
|
|
self.target_size = target_size
|
|
self.cache_dir = Path(cache_dir) if cache_dir is not None else None
|
|
if self.cache_dir is not None:
|
|
self.cache_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
df = pd.read_csv(manifest_path)
|
|
self.entries: Dict[str, dict] = {}
|
|
for _, row in df.iterrows():
|
|
img_path = Path(row["image_path"]).resolve()
|
|
self.entries[str(img_path)] = {
|
|
"annotation_disc": row.get("annotation_disc"),
|
|
"annotation_cup": row.get("annotation_cup"),
|
|
"annotation_type_disc": row.get("annotation_type_disc"),
|
|
"annotation_type_cup": row.get("annotation_type_cup"),
|
|
}
|
|
|
|
def _cache_path(self, image_path: Path) -> Optional[Path]:
|
|
if self.cache_dir is None:
|
|
return None
|
|
return self.cache_dir / f"{image_path.stem}_s{int(self.scale * 100)}.npz"
|
|
|
|
@staticmethod
|
|
def _load_contour(path: Path) -> np.ndarray:
|
|
coords = np.loadtxt(path)
|
|
if coords.ndim == 1:
|
|
coords = coords.reshape(-1, 2)
|
|
return coords
|
|
|
|
@staticmethod
|
|
def _contour_to_mask(coords: np.ndarray, size: tuple[int, int]) -> np.ndarray:
|
|
if coords is None or coords.size == 0:
|
|
return np.zeros((size[1], size[0]), dtype=np.uint8)
|
|
img = Image.new("L", size, 0)
|
|
draw = ImageDraw.Draw(img)
|
|
points = [tuple(map(float, pt)) for pt in coords]
|
|
draw.polygon(points, outline=1, fill=1)
|
|
return np.array(img, dtype=np.uint8)
|
|
|
|
def _load_masks(self, entry: dict, image: Image.Image) -> Optional[Tuple[np.ndarray, np.ndarray]]:
|
|
disc_path = entry.get("annotation_disc")
|
|
cup_path = entry.get("annotation_cup")
|
|
disc_type = (entry.get("annotation_type_disc") or "").lower()
|
|
cup_type = (entry.get("annotation_type_cup") or "").lower()
|
|
|
|
disc_mask: Optional[np.ndarray] = None
|
|
cup_mask: Optional[np.ndarray] = None
|
|
|
|
if disc_path and not pd.isna(disc_path):
|
|
disc_path = Path(disc_path)
|
|
try:
|
|
if disc_type == "mask":
|
|
mask_img = Image.open(disc_path)
|
|
mask_img = mask_img.resize(image.size, Image.NEAREST)
|
|
disc_mask, cup_from_mask = disc_cup_from_mask_image(mask_img)
|
|
if cup_from_mask.sum() > 0:
|
|
cup_mask = cup_from_mask
|
|
elif disc_type == "contour":
|
|
coords = self._load_contour(disc_path)
|
|
disc_mask = self._contour_to_mask(coords, image.size)
|
|
except Exception:
|
|
disc_mask = None
|
|
|
|
if cup_mask is None and cup_path and not pd.isna(cup_path):
|
|
cup_path = Path(cup_path)
|
|
try:
|
|
if cup_type == "mask":
|
|
mask_img = Image.open(cup_path)
|
|
mask_img = mask_img.resize(image.size, Image.NEAREST)
|
|
_, cup_mask = disc_cup_from_mask_image(mask_img)
|
|
elif cup_type == "contour":
|
|
coords = self._load_contour(cup_path)
|
|
cup_mask = self._contour_to_mask(coords, image.size)
|
|
except Exception:
|
|
cup_mask = None
|
|
|
|
if disc_mask is None:
|
|
return None
|
|
disc_mask = (disc_mask > 0).astype(np.uint8)
|
|
if cup_mask is None:
|
|
cup_mask = np.zeros_like(disc_mask, dtype=np.uint8)
|
|
cup_mask = ((cup_mask > 0) & (disc_mask > 0)).astype(np.uint8)
|
|
return disc_mask, cup_mask
|
|
|
|
def _compute_crop_info(self, image: Image.Image, image_path: Path) -> Optional[dict]:
|
|
image_path = Path(image_path).resolve()
|
|
entry = self.entries.get(str(image_path))
|
|
if entry is None:
|
|
return None
|
|
cache_path = self._cache_path(image_path)
|
|
cached_bounds = None
|
|
if cache_path is not None and cache_path.exists():
|
|
data = np.load(cache_path, allow_pickle=False)
|
|
try:
|
|
cached_bounds = {
|
|
"left": float(data["left"]),
|
|
"upper": float(data["upper"]),
|
|
"right": float(data["right"]),
|
|
"lower": float(data["lower"]),
|
|
}
|
|
if "features" in data.files:
|
|
cached_bounds["features"] = data["features"].astype(np.float32)
|
|
return cached_bounds
|
|
except KeyError:
|
|
cached_bounds = None
|
|
|
|
masks = self._load_masks(entry, image)
|
|
if masks is None:
|
|
return cached_bounds
|
|
disc_mask, cup_mask = masks
|
|
try:
|
|
geom = _geometry_from_mask(disc_mask, self.scale)
|
|
except Exception:
|
|
return cached_bounds
|
|
cx = geom["centre_x"]
|
|
cy = geom["centre_y"]
|
|
r = geom["crop_radius"]
|
|
left = max(0.0, cx - r)
|
|
upper = max(0.0, cy - r)
|
|
right = min(float(image.width), cx + r)
|
|
lower = min(float(image.height), cy + r)
|
|
features = compute_geometry_features(disc_mask, cup_mask)
|
|
|
|
info = {
|
|
"left": left,
|
|
"upper": upper,
|
|
"right": right,
|
|
"lower": lower,
|
|
"features": features,
|
|
}
|
|
if cache_path is not None:
|
|
np.savez(
|
|
cache_path,
|
|
left=left,
|
|
upper=upper,
|
|
right=right,
|
|
lower=lower,
|
|
width=float(image.width),
|
|
height=float(image.height),
|
|
scale=self.scale,
|
|
target_size=self.target_size,
|
|
features=features,
|
|
)
|
|
return info
|
|
|
|
def __call__(self, image: Image.Image, image_path: Path) -> Image.Image:
|
|
info = self._compute_crop_info(image, image_path)
|
|
if info is None:
|
|
return image
|
|
left = info["left"]
|
|
upper = info["upper"]
|
|
right = info["right"]
|
|
lower = info["lower"]
|
|
|
|
if right <= left or lower <= upper:
|
|
return image
|
|
crop = image.crop((left, upper, right, lower))
|
|
return crop.resize((self.target_size, self.target_size), Image.BILINEAR)
|
|
|
|
def geometry_features(self, image: Image.Image, image_path: Path) -> Optional[np.ndarray]:
|
|
info = self._compute_crop_info(image, image_path)
|
|
if info is None:
|
|
return None
|
|
features = info.get("features")
|
|
if features is None:
|
|
return None
|
|
return np.asarray(features, dtype=np.float32)
|
|
|
|
def _num(x):
|
|
"""float(x) or None if NaN/None/invalid."""
|
|
try:
|
|
v = float(x)
|
|
except Exception:
|
|
return None
|
|
return None if math.isnan(v) else v
|
|
|
|
class _ClinicalView:
|
|
"""Minimal shim so ClinicalDataset can iterate an epoch-specific DataFrame
|
|
while still delegating encoding/paths/labels to the ClinicalData object."""
|
|
def __init__(self, base: ClinicalData, df):
|
|
self.base = base
|
|
self.df = df
|
|
|
|
@property
|
|
def image_dir(self):
|
|
return self.base.image_dir
|
|
|
|
@property
|
|
def clinical_dir(self):
|
|
return self.base.clinical_dir
|
|
|
|
@property
|
|
def id_cols(self):
|
|
return ("Patient ID", "eyeID")
|
|
|
|
@property
|
|
def label_col(self):
|
|
return self.base.label_col
|
|
|
|
@property
|
|
def filename_template(self):
|
|
# prefer whatever the dataset defined; otherwise fall back to RET{pid}{eye}.jpg style
|
|
return getattr(self.base, "filename_template", "RET{pid:03d}{eye}.jpg")
|
|
|
|
@property
|
|
def dim(self):
|
|
return self.base.feature_dim
|
|
|
|
def encode_metadata(self, row):
|
|
# ClinicalData returns numpy; convert to torch here to keep towers torch-only
|
|
import torch as _torch
|
|
vec = self.base.vectorize_row(row)
|
|
return _torch.as_tensor(vec, dtype=_torch.float32)
|
|
|
|
def get_image_path(self, row):
|
|
return self.base.get_image_path(row)
|
|
|
|
def get_label(self, row):
|
|
return int(row[self.base.label_col])
|
|
|
|
|
|
class HyperTower:
|
|
|
|
|
|
def __init__(self, clinical: ClinicalData, args):
|
|
# Expects a fully built ClinicalData (add_df handled upstream)
|
|
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
print(f"Using device: {self.device}")
|
|
|
|
self.clinical = clinical
|
|
self.warmup_tower_epochs = getattr(args, "warmup_tower_epochs", 2)
|
|
self.warmup_fused_epochs = getattr(args, "warmup_fused_epochs", 3)
|
|
self.args = args # cache for checkpointing
|
|
self._early = None
|
|
self.aux_img = getattr(args, "aux_img", 0.05)
|
|
self.aux_md = getattr(args, "aux_md", 0.05)
|
|
self.aux_detach = getattr(args, "aux_detach", True)
|
|
# SE placement configuration
|
|
se_where = getattr(args, "se_where", "bridge") # 'bridge'|'tower'|'both'|'none'
|
|
use_se_tower = se_where in ("tower", "both")
|
|
# Optional disc-centric cropping
|
|
self.image_preprocessor = None
|
|
crop_manifest = getattr(args, "img_crop_manifest", None)
|
|
crop_weights = getattr(args, "img_crop_weights", None)
|
|
use_gt = getattr(args, "img_crop_gt", False)
|
|
if crop_manifest:
|
|
crop_cache = getattr(args, "img_crop_cache", Path("analysis_data/hypertower_crops"))
|
|
crop_cache = Path(crop_cache)
|
|
if use_gt:
|
|
self.image_preprocessor = ManifestImageCropper(
|
|
manifest_path=Path(crop_manifest),
|
|
scale=getattr(args, "img_crop_scale", 2.5),
|
|
target_size=getattr(args, "img_crop_size", 224),
|
|
cache_dir=crop_cache,
|
|
)
|
|
print(f"[HyperTower] GT disc cropper enabled → cache at {crop_cache}")
|
|
elif crop_weights:
|
|
self.image_preprocessor = UNetImageCropper(
|
|
manifest_path=Path(crop_manifest),
|
|
weights_path=Path(crop_weights),
|
|
normalize=getattr(args, "img_crop_normalize", "per_image"),
|
|
threshold=getattr(args, "img_crop_threshold", 0.5),
|
|
tta=getattr(args, "img_crop_tta", False),
|
|
scale=getattr(args, "img_crop_scale", 2.5),
|
|
target_size=getattr(args, "img_crop_size", 224),
|
|
cache_dir=crop_cache,
|
|
)
|
|
print(f"[HyperTower] UNet disc cropper enabled → cache at {crop_cache}")
|
|
else:
|
|
print("[HyperTower] img_crop_manifest provided but no weights/gt flag; skipping cropping")
|
|
|
|
self.use_geometry_features = bool(getattr(args, "img_geometry_features", False))
|
|
if self.use_geometry_features and self.image_preprocessor is None:
|
|
raise ValueError("img_geometry_features requires --img-crop-manifest with either --img-crop-weights or --img-crop-gt.")
|
|
self.geometry_dim = FEATURE_DIM if self.use_geometry_features else 0
|
|
|
|
# Towers derive their dimensions from ClinicalData
|
|
self.img_tower = ImageTower(
|
|
backbone=getattr(args, "backbone", "efficientnet_b0"),
|
|
freeze_ratio=getattr(args, "freeze_ratio", 0.0),
|
|
use_se=use_se_tower,
|
|
se_reduction=getattr(args, "se_reduction_tower", getattr(args, "se_reduction", 16)),
|
|
se_pre_norm=getattr(args, "se_pre_norm_tower", getattr(args, "se_pre_norm", True)),
|
|
augment=getattr(args, "img_augment", True),
|
|
geometry_dim=self.geometry_dim,
|
|
).to(self.device)
|
|
self.md_tower = MDTower(
|
|
self.clinical,
|
|
use_se=use_se_tower,
|
|
se_reduction=getattr(args, "se_reduction_tower", getattr(args, "se_reduction", 16)),
|
|
se_pre_norm=getattr(args, "se_pre_norm_tower", getattr(args, "se_pre_norm", True)),
|
|
).to(self.device)
|
|
self.mode = args.fusion_mode # cache
|
|
# Training hyperparams / objects
|
|
self.batch_size = args.batch_size
|
|
self.epochs = args.epochs
|
|
self.lr = args.lr
|
|
self.fold = args.fold
|
|
self.criterion = nn.CrossEntropyLoss()
|
|
self._ce_weight = self.criterion.weight.detach().clone() if self.criterion.weight is not None else None
|
|
if self._ce_weight is not None:
|
|
self._ce_weight = self._ce_weight.to(self.device)
|
|
self._ce_reduction = self.criterion.reduction
|
|
self.focal_gamma = float(getattr(args, "focal_gamma", 0.0) or 0.0)
|
|
# Run and model directories come from the script
|
|
self.run_dir = Path(getattr(args, "run_dir", ".")).resolve()
|
|
self.run_dir.mkdir(parents=True, exist_ok=True)
|
|
# Per-run log locations to avoid collisions across concurrent workers
|
|
self.epoch_log_path = self.run_dir / "epoch_log.csv"
|
|
self.train_log_path = self.run_dir / "train.log"
|
|
self.models_dir = Path(getattr(args, "models_dir", "models")).resolve()
|
|
self.models_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
self.holdout_df = getattr(args, "holdout_df", None)
|
|
self.holdout_loader = None
|
|
if isinstance(self.holdout_df, pd.DataFrame) and not self.holdout_df.empty:
|
|
if getattr(args, "eval_mode", "multiclass") == "binary":
|
|
label_col = getattr(self.clinical, "label_col", None)
|
|
if label_col and label_col in self.holdout_df.columns:
|
|
self.holdout_df = self.holdout_df[self.holdout_df[label_col].isin([0, 1])].reset_index(drop=True)
|
|
self.holdout_loader = self._make_loader_for_df(self.holdout_df, is_train=False)
|
|
print(f"[HyperTower] Holdout loader prepared with {len(self.holdout_df)} samples")
|
|
|
|
if self.mode == "vote":
|
|
# Per-tower classification heads (logits) + vote combiner
|
|
self.head_img = nn.Linear(self.img_tower.out_dim, args.num_classes).to(self.device)
|
|
self.head_md = nn.Linear(self.md_tower.out_dim, args.num_classes).to(self.device)
|
|
self.vote = VoteBridge(num_classes=args.num_classes).to(self.device)
|
|
|
|
# Optimizer: towers + heads + vote
|
|
self.optimizer = torch.optim.Adam(
|
|
list(self.img_tower.parameters())
|
|
+ list(self.md_tower.parameters())
|
|
+ list(self.head_img.parameters())
|
|
+ list(self.head_md.parameters())
|
|
+ list(self.vote.parameters()),
|
|
lr=self.lr,
|
|
)
|
|
else:
|
|
# Existing feature-fusion bridge
|
|
self.bridge = Bridge(
|
|
img_dim=self.img_tower.out_dim,
|
|
meta_dim=self.md_tower.out_dim,
|
|
num_classes=args.num_classes,
|
|
fusion_dim=256,
|
|
mode=args.fusion_mode,
|
|
use_se=(se_where in ("bridge","both")) and getattr(args, "use_se", True),
|
|
se_reduction=(getattr(args, "se_reduction", 16) if getattr(args, "use_se", True) else 0),
|
|
se_pre_norm=getattr(args, "se_pre_norm", True)
|
|
).to(self.device)
|
|
|
|
self.optimizer = torch.optim.Adam(
|
|
list(self.img_tower.parameters())
|
|
+ list(self.md_tower.parameters())
|
|
+ list(self.bridge.parameters()),
|
|
lr=self.lr,
|
|
)
|
|
|
|
|
|
# EMA-forgiveness + BCD knobs; logging stays simple for now
|
|
self.ema_alpha = args.ema_alpha
|
|
self.bcd_prob = args.bcd_prob
|
|
self.ema_fused_loss = None
|
|
|
|
# File logger (per-run); simple line format, no duplication to root
|
|
self.logger = logging.getLogger("hypertower")
|
|
self.logger.setLevel(logging.INFO)
|
|
# Replace existing handlers to avoid duplicate lines across folds
|
|
self.logger.handlers = []
|
|
fh = logging.FileHandler(str(self.train_log_path))
|
|
fh.setFormatter(logging.Formatter("%(asctime)s - %(message)s"))
|
|
self.logger.addHandler(fh)
|
|
self.logger.propagate = False
|
|
|
|
# Dynamic BCD configuration (epoch baselines + batch nudges)
|
|
self.bcd_cfg = {
|
|
"metric": getattr(args, "bcd_metric", "auc"),
|
|
"p0": getattr(args, "bcd_p0", 0.20),
|
|
"k": getattr(args, "bcd_k", 0.4),
|
|
"pmin": getattr(args, "bcd_min", 0.05),
|
|
"pmax": getattr(args, "bcd_max", 0.30),
|
|
"alpha_batch": getattr(args, "bcd_alpha_batch", 0.2),
|
|
"alpha_tower": getattr(args, "bcd_alpha_tower", 0.3),
|
|
"explore_floor": getattr(args, "bcd_explore_floor", 0.15),
|
|
"entropy_ema": getattr(args, "entropy_ema", 0.7),
|
|
}
|
|
|
|
# Track last epoch's eval metrics for baseline deficits; initialize safely
|
|
self.last_eval = {
|
|
"acc_fused": 0.0, "acc_img": 0.0, "acc_md": 0.0,
|
|
"auc_fused": 0.0, "auc_img": 0.0, "auc_md": 0.0,
|
|
}
|
|
# Running EMA of per-head uncertainty (entropy in [0,1])
|
|
self.entropy_ema = {"fused": 0.5, "img": 0.5, "md": 0.5}
|
|
|
|
# DataLoaders are (re)built each epoch from ClinicalData's splits
|
|
self.train_loader = None
|
|
self.test_loader = None
|
|
self._last_step_mix = None
|
|
|
|
if not getattr(self.clinical, "folds", None):
|
|
raise RuntimeError("ClinicalData has no built folds. Did you call add_df(...) upstream?")
|
|
|
|
@staticmethod
|
|
def _confusion(preds, labels):
|
|
# Quick TP/TN/FP/FN for binary debug (kept as-is)
|
|
tp = ((preds == 1) & (labels == 1)).sum().item()
|
|
tn = ((preds == 0) & (labels == 0)).sum().item()
|
|
fp = ((preds == 1) & (labels == 0)).sum().item()
|
|
fn = ((preds == 0) & (labels == 1)).sum().item()
|
|
return {"tp": tp, "tn": tn, "fp": fp, "fn": fn}
|
|
|
|
def _unpack_batch(self, batch):
|
|
if len(batch) == 4:
|
|
imgs, metas, geometry, labels = batch
|
|
else:
|
|
imgs, metas, labels = batch
|
|
if self.geometry_dim > 0:
|
|
geometry = torch.zeros(imgs.size(0), self.geometry_dim, dtype=torch.float32)
|
|
else:
|
|
geometry = None
|
|
return imgs, metas, geometry, labels
|
|
|
|
def _classification_loss(self, logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
|
|
"""
|
|
Shared classification loss (CrossEntropy or focal depending on gamma).
|
|
"""
|
|
weight = self._ce_weight
|
|
if weight is not None and weight.device != logits.device:
|
|
weight = weight.to(logits.device)
|
|
return focal_loss(
|
|
logits,
|
|
labels,
|
|
gamma=self.focal_gamma,
|
|
weight=weight,
|
|
reduction=self._ce_reduction,
|
|
)
|
|
|
|
def _step_image_only(self, imgs, metas, geometry, labels):
|
|
if self.mode == "vote":
|
|
img_feats = self.img_tower(imgs, geometry)
|
|
out_img = self.head_img(img_feats)
|
|
loss_i = self._classification_loss(out_img, labels)
|
|
preds_i = out_img.argmax(dim=1)
|
|
acc_i = (preds_i == labels).float().mean().item()
|
|
cm = self._confusion(preds_i, labels)
|
|
|
|
self.optimizer.zero_grad()
|
|
loss_i.backward()
|
|
self.optimizer.step()
|
|
return None, loss_i.item(), None, None, acc_i, None, {"img": cm}
|
|
else:# Optimize the image head alone (used when BCD chooses image-only)
|
|
img_feats = self.img_tower(imgs, geometry)
|
|
out_img = self.bridge.classifier_img(img_feats)
|
|
loss_i = self._classification_loss(out_img, labels)
|
|
preds_i = out_img.argmax(dim=1)
|
|
acc_i = (preds_i == labels).float().mean().item()
|
|
cm = self._confusion(preds_i, labels)
|
|
self.optimizer.zero_grad()
|
|
loss_i.backward()
|
|
self.optimizer.step()
|
|
return None, loss_i.item(), None, None, acc_i, None, {"img": cm}
|
|
|
|
def _step_meta_only(self, imgs, metas, geometry, labels):
|
|
# Optimize the metadata head alone (used when BCD chooses meta-only)
|
|
if self.mode == "vote":
|
|
meta_feats = self.md_tower(metas)
|
|
out_md = self.head_md(meta_feats)
|
|
loss_m = self._classification_loss(out_md, labels)
|
|
preds_m = out_md.argmax(dim=1)
|
|
acc_m = (preds_m == labels).float().mean().item()
|
|
cm = self._confusion(preds_m, labels)
|
|
|
|
self.optimizer.zero_grad()
|
|
loss_m.backward()
|
|
self.optimizer.step()
|
|
return None, None, loss_m.item(), None, None, acc_m, {"meta": cm}
|
|
else:
|
|
meta_feats = self.md_tower(metas)
|
|
out_md = self.bridge.classifier_md(meta_feats)
|
|
loss_m = self._classification_loss(out_md, labels)
|
|
preds_m = out_md.argmax(dim=1)
|
|
acc_m = (preds_m == labels).float().mean().item()
|
|
cm = self._confusion(preds_m, labels)
|
|
self.optimizer.zero_grad()
|
|
loss_m.backward()
|
|
self.optimizer.step()
|
|
return None, None, loss_m.item(), None, None, acc_m, {"meta": cm}
|
|
|
|
def _step_fused(self, imgs, metas, geometry, labels):
|
|
if self.mode == "vote":
|
|
img_feats = self.img_tower(imgs, geometry)
|
|
md_feats = self.md_tower(metas)
|
|
|
|
out_img = self.head_img(img_feats)
|
|
out_md = self.head_md(md_feats)
|
|
out_fused = self.vote(out_img, out_md)
|
|
|
|
preds_f = out_fused.argmax(dim=1)
|
|
acc_f = (preds_f == labels).float().mean().item()
|
|
loss_f = self._classification_loss(out_fused, labels)
|
|
loss_f_val = float(loss_f.detach().item())
|
|
|
|
# Optional: keep your EMA + small aux tower losses just like before
|
|
if self.ema_fused_loss is None:
|
|
self.ema_fused_loss = loss_f_val
|
|
L_total = (1.0 - self.ema_alpha) * loss_f + self.ema_alpha * loss_f.new_tensor(self.ema_fused_loss)
|
|
|
|
aux_terms = 0.0
|
|
if self.aux_img > 0:
|
|
aux_terms = aux_terms + self.aux_img * self._classification_loss(
|
|
out_img.detach() if self.aux_detach else out_img,
|
|
labels,
|
|
)
|
|
if self.aux_md > 0:
|
|
aux_terms = aux_terms + self.aux_md * self._classification_loss(
|
|
out_md.detach() if self.aux_detach else out_md,
|
|
labels,
|
|
)
|
|
L_total = L_total + aux_terms
|
|
|
|
self.optimizer.zero_grad()
|
|
L_total.backward()
|
|
self.optimizer.step()
|
|
|
|
self.ema_fused_loss = self.ema_alpha * self.ema_fused_loss + (1.0 - self.ema_alpha) * loss_f_val
|
|
|
|
# Confusion tables for logging
|
|
cm = {
|
|
"fused": self._confusion(preds_f, labels),
|
|
"img": self._confusion(out_img.argmax(dim=1), labels),
|
|
"meta": self._confusion(out_md.argmax(dim=1), labels),
|
|
}
|
|
acc_i = (out_img.argmax(dim=1) == labels).float().mean().item()
|
|
acc_m = (out_md.argmax(dim=1) == labels).float().mean().item()
|
|
|
|
with torch.no_grad():
|
|
loss_img_val = self._classification_loss(out_img, labels).item()
|
|
loss_md_val = self._classification_loss(out_md, labels).item()
|
|
return loss_f_val, loss_img_val, loss_md_val, acc_f, acc_i, acc_m, cm
|
|
else:
|
|
img_feats = self.img_tower(imgs, geometry)
|
|
meta_feats = self.md_tower(metas)
|
|
out_fused, out_img, out_md = self.bridge(img_feats, meta_feats)
|
|
|
|
preds_f = out_fused.argmax(dim=1)
|
|
acc_f = (preds_f == labels).float().mean().item()
|
|
loss_f = self._classification_loss(out_fused, labels)
|
|
loss_f_val = float(loss_f.detach().item())
|
|
|
|
cm = {"fused": self._confusion(preds_f, labels)}
|
|
loss_i = acc_i = loss_m = acc_m = None
|
|
|
|
if out_img is not None:
|
|
with torch.no_grad():
|
|
preds_i = out_img.argmax(dim=1)
|
|
acc_i = (preds_i == labels).float().mean().item()
|
|
loss_i = self._classification_loss(out_img, labels).item()
|
|
cm["img"] = self._confusion(preds_i, labels)
|
|
if out_md is not None:
|
|
with torch.no_grad():
|
|
preds_m = out_md.argmax(dim=1)
|
|
acc_m = (preds_m == labels).float().mean().item()
|
|
loss_m = self._classification_loss(out_md, labels).item()
|
|
cm["meta"] = self._confusion(preds_m, labels)
|
|
|
|
# EMA-blended fused loss
|
|
if self.ema_fused_loss is None:
|
|
self.ema_fused_loss = loss_f_val
|
|
L_total = (1.0 - self.ema_alpha) * loss_f + self.ema_alpha * loss_f.new_tensor(self.ema_fused_loss)
|
|
|
|
# Auxiliary tower losses (small weights). Use detached features to calibrate heads only.
|
|
aux_terms = 0.0
|
|
if self.aux_img > 0 and out_img is not None:
|
|
logits_img_for_aux = self.bridge.classifier_img(img_feats.detach()) if self.aux_detach else out_img
|
|
aux_terms = aux_terms + self.aux_img * self._classification_loss(logits_img_for_aux, labels)
|
|
if self.aux_md > 0 and out_md is not None:
|
|
logits_md_for_aux = self.bridge.classifier_md(meta_feats.detach()) if self.aux_detach else out_md
|
|
aux_terms = aux_terms + self.aux_md * self._classification_loss(logits_md_for_aux, labels)
|
|
|
|
L_total = L_total + aux_terms
|
|
|
|
self.optimizer.zero_grad()
|
|
L_total.backward()
|
|
self.optimizer.step()
|
|
|
|
# update EMA after the step
|
|
self.ema_fused_loss = self.ema_alpha * self.ema_fused_loss + (1.0 - self.ema_alpha) * loss_f_val
|
|
|
|
return loss_f_val, loss_i, loss_m, acc_f, acc_i, acc_m, cm
|
|
|
|
|
|
def _make_loader_for_df(self, df, is_train: bool = False):
|
|
# Rebuild a DataLoader for the current epoch's split
|
|
view = _ClinicalView(self.clinical, df)
|
|
geometry_provider = self.image_preprocessor if self.geometry_dim > 0 else None
|
|
ds = ClinicalDataset(
|
|
view,
|
|
self.img_tower.transform,
|
|
image_preprocessor=self.image_preprocessor,
|
|
geometry_provider=geometry_provider,
|
|
geometry_dim=self.geometry_dim,
|
|
)
|
|
|
|
# Optional: class-balanced bootstrapped sampling for TRAIN only
|
|
if is_train and getattr(self.args, "balanced_sampler", False):
|
|
import numpy as _np
|
|
y = _np.asarray(df[self.clinical.label_col].values)
|
|
# inverse-frequency weights per class
|
|
uniq, counts = _np.unique(y, return_counts=True)
|
|
inv = {c: (1.0 / cnt if cnt > 0 else 0.0) for c, cnt in zip(uniq, counts)}
|
|
w = _np.array([inv[c] for c in y], dtype=_np.float32)
|
|
sampler = WeightedRandomSampler(weights=w, num_samples=len(y), replacement=True)
|
|
return DataLoader(ds, batch_size=self.batch_size, sampler=sampler, shuffle=False)
|
|
|
|
return DataLoader(ds, batch_size=self.batch_size, shuffle=not is_train)
|
|
|
|
# ---- Dynamic BCD helpers ----
|
|
def _metric_value(self, name: str) -> float:
|
|
# Pull either AUC or ACC from last_eval, falling back if NaN/zero
|
|
if self.bcd_cfg["metric"] == "auc":
|
|
v = self.last_eval.get(f"auc_{name}", 0.0)
|
|
if v == v: # not NaN
|
|
return float(v)
|
|
# fallback to accuracy
|
|
return float(self.last_eval.get(f"acc_{name}", 0.0))
|
|
return float(self.last_eval.get(f"acc_{name}", 0.0))
|
|
|
|
def _epoch_bcd_baseline(self):
|
|
# Compute epoch-level baselines: p_bcd_epoch and tower weights w_i, w_m
|
|
p0 = self.bcd_cfg["p0"]; k = self.bcd_cfg["k"]
|
|
pmin = self.bcd_cfg["pmin"]; pmax = self.bcd_cfg["pmax"]
|
|
eps = 1e-6
|
|
Af = self._metric_value("fused"); Ai = self._metric_value("img"); Am = self._metric_value("md")
|
|
di = max(0.0, Af - Ai); dm = max(0.0, Af - Am)
|
|
p_bcd_epoch = max(pmin, min(pmax, p0 + k * (di + dm) / 2.0))
|
|
wi = (di + eps) / (di + dm + 2 * eps)
|
|
wm = 1.0 - wi
|
|
return p_bcd_epoch, wi, wm, {"di": di, "dm": dm, "Af": Af, "Ai": Ai, "Am": Am}
|
|
|
|
@staticmethod
|
|
def _entropy_from_logits(logits: torch.Tensor, num_classes: int) -> float:
|
|
# Returns entropy normalized to [0,1] using log(K) denominator
|
|
with torch.no_grad():
|
|
probs = F.softmax(logits, dim=1)
|
|
ent = -(probs * (probs.clamp_min(1e-12)).log()).sum(dim=1)
|
|
ent = ent / np.log(num_classes)
|
|
return float(ent.mean().item())
|
|
|
|
def _batch_uncertainty(self, imgs: torch.Tensor, metas: torch.Tensor, geometry: Optional[torch.Tensor] = None) -> dict:
|
|
self.img_tower.eval(); self.md_tower.eval()
|
|
if self.mode == "vote":
|
|
self.head_img.eval(); self.head_md.eval(); self.vote.eval()
|
|
with torch.no_grad():
|
|
if geometry is None and self.geometry_dim > 0:
|
|
geometry = torch.zeros(imgs.size(0), self.geometry_dim, device=imgs.device, dtype=imgs.dtype)
|
|
img_feats = self.img_tower(imgs, geometry)
|
|
md_feats = self.md_tower(metas)
|
|
out_img = self.head_img(img_feats)
|
|
out_md = self.head_md(md_feats)
|
|
out_fused = self.vote(out_img, out_md)
|
|
K = out_fused.shape[1]
|
|
e_f = self._entropy_from_logits(out_fused, K)
|
|
e_i = self._entropy_from_logits(out_img, K)
|
|
e_m = self._entropy_from_logits(out_md, K)
|
|
# restore train()
|
|
self.img_tower.train(); self.md_tower.train()
|
|
self.head_img.train(); self.head_md.train(); self.vote.train()
|
|
else:
|
|
self.bridge.eval()
|
|
with torch.no_grad():
|
|
if geometry is None and self.geometry_dim > 0:
|
|
geometry = torch.zeros(imgs.size(0), self.geometry_dim, device=imgs.device, dtype=imgs.dtype)
|
|
img_feats = self.img_tower(imgs, geometry)
|
|
meta_feats = self.md_tower(metas)
|
|
out_fused, out_img, out_md = self.bridge(img_feats, meta_feats)
|
|
K = out_fused.shape[1]
|
|
e_f = self._entropy_from_logits(out_fused, K)
|
|
e_i = self._entropy_from_logits(out_img, K) if out_img is not None else 0.5
|
|
e_m = self._entropy_from_logits(out_md, K) if out_md is not None else 0.5
|
|
self.img_tower.train(); self.md_tower.train(); self.bridge.train()
|
|
|
|
a = self.bcd_cfg["entropy_ema"]
|
|
self.entropy_ema["fused"] = a * self.entropy_ema["fused"] + (1 - a) * e_f
|
|
self.entropy_ema["img"] = a * self.entropy_ema["img"] + (1 - a) * e_i
|
|
self.entropy_ema["md"] = a * self.entropy_ema["md"] + (1 - a) * e_m
|
|
return {"fused": self.entropy_ema["fused"], "img": self.entropy_ema["img"], "md": self.entropy_ema["md"]}
|
|
|
|
def _write_epoch_log(self, row: dict, path: str | Path | None = None):
|
|
"""
|
|
Write an epoch row. Initializes the CSV with a stable header on first call.
|
|
Later calls auto-fill missing columns with None so nothing gets dropped.
|
|
"""
|
|
# known optional columns we might add later
|
|
OPTIONAL_COLS = [
|
|
"pct_fused", "pct_img", "pct_md", "phase",
|
|
"se_mean", "se_std", "se_pct_lt_0.2", "se_pct_gt_0.8",
|
|
"holdout_loss", "holdout_acc_fused", "holdout_acc_img", "holdout_acc_md",
|
|
"holdout_auc_fused", "holdout_auc_img", "holdout_auc_md",
|
|
]
|
|
|
|
# On first call: create file, lock fieldnames
|
|
if not hasattr(self, "_epoch_log_writer"):
|
|
# union of current row keys + optional columns so header includes them even if None now
|
|
fieldnames = list(dict.fromkeys([*row.keys(), *OPTIONAL_COLS]))
|
|
self._epoch_log_path = Path(path) if path is not None else self.epoch_log_path
|
|
self._epoch_log_path.parent.mkdir(parents=True, exist_ok=True)
|
|
self._epoch_log_fp = open(self._epoch_log_path, "w", newline="")
|
|
self._epoch_log_writer = csv.DictWriter(self._epoch_log_fp, fieldnames=fieldnames)
|
|
self._epoch_log_writer.writeheader()
|
|
self._epoch_log_fields = fieldnames # remember for future rows
|
|
|
|
# Ensure all header fields exist in this row
|
|
for k in self._epoch_log_fields:
|
|
row.setdefault(k, None)
|
|
|
|
# Write and flush
|
|
self._epoch_log_writer.writerow({k: row.get(k) for k in self._epoch_log_fields})
|
|
self._epoch_log_fp.flush()
|
|
|
|
def _snapshot(self):
|
|
state = {
|
|
"img_tower": self.img_tower.state_dict(),
|
|
"md_tower": self.md_tower.state_dict(),
|
|
"optimizer": self.optimizer.state_dict(),
|
|
"args": vars(self.args),
|
|
}
|
|
if hasattr(self, "bridge"): state["bridge"] = self.bridge.state_dict()
|
|
if hasattr(self, "head_img"): state["head_img"] = self.head_img.state_dict()
|
|
if hasattr(self, "head_md"): state["head_md"] = self.head_md.state_dict()
|
|
return state
|
|
|
|
|
|
def train(self):
|
|
best_path = None
|
|
if getattr(self.args, "checkpoint_best", False):
|
|
best_path = str((self.models_dir / "best.pth").resolve())
|
|
|
|
# Always track the best model even if we don't stop early
|
|
monitor = getattr(self.args, "early_metric", None)
|
|
if not monitor:
|
|
# Prefer AUC-based monitoring by default; fall back to loss only if specified
|
|
if self.mode == "image_only":
|
|
monitor = "auc_img"
|
|
elif self.mode == "metadata_only":
|
|
monitor = "auc_md"
|
|
else: # fused or vote
|
|
monitor = "auc_fused"
|
|
mode = getattr(self.args, "early_mode", "auto")
|
|
# Optional switch: monitor holdout metrics instead of validation
|
|
use_holdout_monitor = bool(getattr(self.args, "early_monitor_holdout", False))
|
|
if use_holdout_monitor and self.holdout_loader is None:
|
|
print("[early] Requested holdout monitoring but no holdout set is configured; falling back to validation metrics.")
|
|
use_holdout_monitor = False
|
|
if use_holdout_monitor:
|
|
# If the monitor isn't already a holdout metric, prepend it
|
|
if not monitor.startswith("holdout_"):
|
|
monitor = f"holdout_{monitor}"
|
|
# Force sensible default if no monitor was provided
|
|
if monitor == "holdout_auc_fused" and mode == "auto":
|
|
mode = "max"
|
|
if mode == "auto":
|
|
mode = ("min" if ("loss" in monitor.lower()) else "max")
|
|
self._best_info = {
|
|
"monitor": monitor,
|
|
"mode": mode,
|
|
"min_delta": float(getattr(self.args, "early_min_delta", 0.0)),
|
|
"value": (-math.inf if mode == "max" else math.inf),
|
|
"epoch": -1,
|
|
"state": None,
|
|
"save_path": best_path,
|
|
}
|
|
# keep args in sync so EarlyStopper uses the same monitor/mode
|
|
self.args.early_metric = monitor
|
|
self.args.early_mode = mode
|
|
# Separate tracker for best holdout performance (if a holdout set is provided)
|
|
self._holdout_best = None
|
|
if self.holdout_loader is not None:
|
|
self._holdout_best = {
|
|
"monitor": "holdout_auc_fused",
|
|
"mode": "max",
|
|
"min_delta": 0.0,
|
|
"value": -math.inf,
|
|
"epoch": -1,
|
|
"state": None,
|
|
"save_path": str((self.models_dir / "model_holdout_best.pt").resolve()),
|
|
}
|
|
|
|
self._early = None
|
|
if getattr(self.args, "early_stop", False):
|
|
self._early = EarlyStopper(
|
|
monitor=monitor,
|
|
mode=mode,
|
|
patience=self.args.early_patience,
|
|
min_delta=self.args.early_min_delta,
|
|
save_path=best_path,
|
|
restore_best=True,
|
|
)
|
|
# Gradual thaw config
|
|
gradual = bool(getattr(self.args, "gradual_thaw", False))
|
|
thaw_ratio = float(getattr(self.args, "thaw_ratio", 0.33))
|
|
thaw_phase = int(getattr(self.args, "thaw_phase_duration", 5))
|
|
thaw_target = str(getattr(self.args, "thaw_target", "image"))
|
|
thaw_start_epoch = int(getattr(self.args, "thaw_start_epoch", -1))
|
|
thaw_initial_freeze = bool(getattr(self.args, "thaw_initial_freeze", False))
|
|
if thaw_start_epoch < 0:
|
|
thaw_start_epoch = int(getattr(self, "warmup_tower_epochs", 0))
|
|
|
|
for epoch in range(self.epochs):
|
|
# Apply thaw schedule at epoch start: start fully frozen and unfreeze by ratios
|
|
if gradual and thaw_phase > 0:
|
|
try:
|
|
if epoch < thaw_start_epoch:
|
|
if thaw_initial_freeze:
|
|
freeze_r = 1.0
|
|
if thaw_target in ("image", "both") and hasattr(self, "img_tower"):
|
|
self.img_tower.set_freeze_ratio(freeze_r)
|
|
if thaw_target in ("metadata", "both") and hasattr(self, "md_tower"):
|
|
self.md_tower.set_freeze_ratio(freeze_r)
|
|
# else: leave current requires_grad as constructed (no forced freeze)
|
|
else:
|
|
phase_idx = (epoch - thaw_start_epoch) // thaw_phase
|
|
freeze_r = max(0.0, 1.0 - phase_idx * thaw_ratio)
|
|
if thaw_target in ("image", "both") and hasattr(self, "img_tower"):
|
|
self.img_tower.set_freeze_ratio(freeze_r)
|
|
if thaw_target in ("metadata", "both") and hasattr(self, "md_tower"):
|
|
self.md_tower.set_freeze_ratio(freeze_r)
|
|
except Exception:
|
|
pass
|
|
# build splits/loaders per-epoch
|
|
train_df, test_df = self.clinical.get_split_dfs(self.fold)
|
|
if hasattr(self.bridge, "reset_se_stats"):
|
|
self.bridge.reset_se_stats()
|
|
if getattr(self, "args", None) is None:
|
|
self.args = argparse.Namespace() # if not stashed already
|
|
# Expect user to pass --eval_mode and --num-classes==2 for binary
|
|
if getattr(self.args, "eval_mode", "multiclass") == "binary":
|
|
# Drop Suspect (assumed label==2 in PAPILA spreadsheets)
|
|
train_df = train_df[train_df[self.clinical.label_col].isin([0,1])].copy()
|
|
test_df = test_df [test_df [self.clinical.label_col].isin([0,1])].copy()
|
|
self.train_loader = self._make_loader_for_df(train_df, is_train=True)
|
|
self.test_loader = self._make_loader_for_df(test_df, is_train=False)
|
|
|
|
self.img_tower.train(); self.md_tower.train(); self.bridge.train()
|
|
|
|
total_losses = {"fused":0.0,"image_only":0.0,"metadata_only":0.0}
|
|
total_accs = {"fused":0.0,"image_only":0.0,"metadata_only":0.0}
|
|
counts = {"fused":0, "image_only":0, "metadata_only":0}
|
|
step_counts = {"fused":0, "image_only":0, "metadata_only":0}
|
|
|
|
# phase selection
|
|
in_tower_warmup = epoch < self.warmup_tower_epochs
|
|
in_fused_warmup = (self.warmup_tower_epochs <= epoch < (self.warmup_tower_epochs + self.warmup_fused_epochs))
|
|
|
|
# epoch baselines (needed for splits and logging)
|
|
p_bcd_epoch, wi_epoch, wm_epoch, diag = self._epoch_bcd_baseline()
|
|
|
|
if in_tower_warmup:
|
|
# tower-only warmup: alternate or use epoch split; default to 0.5 if empty
|
|
wi = wi_epoch if (diag["Ai"] or diag["Am"]) else 0.5
|
|
for batch_idx, batch in enumerate(self.train_loader):
|
|
imgs, metas, geometry, labels = self._unpack_batch(batch)
|
|
imgs = imgs.to(self.device)
|
|
metas = metas.to(self.device)
|
|
labels = labels.to(self.device)
|
|
geometry = geometry.to(self.device) if geometry is not None else None
|
|
if random() < wi:
|
|
step_type = "image"
|
|
loss_f, loss_i, loss_m, acc_f, acc_i, acc_m, cm = self._step_image_only(imgs, metas, geometry, labels)
|
|
step_counts["image_only"] += 1
|
|
else:
|
|
step_type = "meta"
|
|
loss_f, loss_i, loss_m, acc_f, acc_i, acc_m, cm = self._step_meta_only(imgs, metas, geometry, labels)
|
|
step_counts["metadata_only"] += 1
|
|
|
|
self.logger.info(f"epoch={epoch} batch={batch_idx} phase=tower_warmup step={step_type} "
|
|
f"loss_f={loss_f} loss_i={loss_i} loss_m={loss_m} acc_f={acc_f} acc_i={acc_i} acc_m={acc_m} cm={cm}")
|
|
|
|
if loss_i is not None:
|
|
total_losses["image_only"] += loss_i; total_accs["image_only"] += acc_i; counts["image_only"] += 1
|
|
if loss_m is not None:
|
|
total_losses["metadata_only"] += loss_m; total_accs["metadata_only"] += acc_m; counts["metadata_only"] += 1
|
|
|
|
elif in_fused_warmup:
|
|
# fused-only warmup
|
|
for batch_idx, batch in enumerate(self.train_loader):
|
|
imgs, metas, geometry, labels = self._unpack_batch(batch)
|
|
imgs = imgs.to(self.device)
|
|
metas = metas.to(self.device)
|
|
labels = labels.to(self.device)
|
|
geometry = geometry.to(self.device) if geometry is not None else None
|
|
step_type = "fused"
|
|
loss_f, loss_i, loss_m, acc_f, acc_i, acc_m, cm = self._step_fused(imgs, metas, geometry, labels)
|
|
step_counts["fused"] += 1
|
|
|
|
self.logger.info(f"epoch={epoch} batch={batch_idx} phase=fused_warmup step={step_type} "
|
|
f"loss_f={loss_f} acc_f={acc_f} cm={cm}")
|
|
|
|
total_losses["fused"] += loss_f; total_accs["fused"] += acc_f; counts["fused"] += 1
|
|
|
|
else:
|
|
# dynamic BCD (fusion-centric)
|
|
self.logger.info(
|
|
f"epoch={epoch} bcd_epoch={p_bcd_epoch:.3f} wi_epoch={wi_epoch:.3f} wm_epoch={wm_epoch:.3f} "
|
|
f"deficits={{img:{diag['di']:.3f}, md:{diag['dm']:.3f}}} metrics={{Af:{diag['Af']:.3f}, Ai:{diag['Ai']:.3f}, Am:{diag['Am']:.3f}}}"
|
|
)
|
|
for batch_idx, batch in enumerate(self.train_loader):
|
|
imgs, metas, geometry, labels = self._unpack_batch(batch)
|
|
imgs = imgs.to(self.device)
|
|
metas = metas.to(self.device)
|
|
labels = labels.to(self.device)
|
|
geometry = geometry.to(self.device) if geometry is not None else None
|
|
ents = self._batch_uncertainty(imgs, metas, geometry)
|
|
|
|
# flip rule: higher fused entropy => fewer tower-only batches => more fused training
|
|
alpha_b = self.bcd_cfg["alpha_batch"]; pmin = self.bcd_cfg["pmin"]; pmax = self.bcd_cfg["pmax"]
|
|
p_bcd_batch = (1 - alpha_b) * p_bcd_epoch + alpha_b * (1.0 - ents["fused"])
|
|
p_bcd_batch = max(pmin, min(pmax, p_bcd_batch))
|
|
|
|
# image vs meta split
|
|
alpha_t = self.bcd_cfg["alpha_tower"]; floor = self.bcd_cfg["explore_floor"]; eps = 1e-6
|
|
s_img = ents["img"] / (ents["img"] + ents["md"] + eps)
|
|
p_img_batch = (1 - alpha_t) * wi_epoch + alpha_t * s_img
|
|
p_img_batch = max(floor, min(1.0 - floor, p_img_batch))
|
|
|
|
u = random(); v = random()
|
|
if u < p_bcd_batch:
|
|
if v < p_img_batch:
|
|
step_type = "image"
|
|
loss_f, loss_i, loss_m, acc_f, acc_i, acc_m, cm = self._step_image_only(imgs, metas, geometry, labels)
|
|
step_counts["image_only"] += 1
|
|
else:
|
|
step_type = "meta"
|
|
loss_f, loss_i, loss_m, acc_f, acc_i, acc_m, cm = self._step_meta_only(imgs, metas, geometry, labels)
|
|
step_counts["metadata_only"] += 1
|
|
else:
|
|
step_type = "fused"
|
|
loss_f, loss_i, loss_m, acc_f, acc_i, acc_m, cm = self._step_fused(imgs, metas, geometry, labels)
|
|
step_counts["fused"] += 1
|
|
|
|
self.logger.info(
|
|
f"epoch={epoch} batch={batch_idx} phase=dynamic step={step_type} "
|
|
f"p_bcd_batch={p_bcd_batch:.3f} p_img_batch={p_img_batch:.3f} "
|
|
f"ents={{f:{ents['fused']:.3f}, i:{ents['img']:.3f}, m:{ents['md']:.3f}}} "
|
|
f"loss_f={loss_f} loss_i={loss_i} loss_m={loss_m} acc_f={acc_f} acc_i={acc_i} acc_m={acc_m} cm={cm}"
|
|
)
|
|
|
|
if loss_f is not None:
|
|
total_losses["fused"] += loss_f; total_accs["fused"] += acc_f; counts["fused"] += 1
|
|
if loss_i is not None:
|
|
total_losses["image_only"] += loss_i; total_accs["image_only"] += acc_i; counts["image_only"] += 1
|
|
if loss_m is not None:
|
|
total_losses["metadata_only"] += loss_m; total_accs["metadata_only"] += acc_m; counts["metadata_only"] += 1
|
|
|
|
# epoch prints
|
|
avg_loss = total_losses["fused"] / counts["fused"] if counts["fused"] else 0.0
|
|
avg_fused = total_accs["fused"] / counts["fused"] if counts["fused"] else 0.0
|
|
avg_img = total_accs["image_only"] / counts["image_only"] if counts["image_only"] else 0.0
|
|
avg_md = total_accs["metadata_only"] / counts["metadata_only"] if counts["metadata_only"] else 0.0
|
|
|
|
# realized mix for logging
|
|
tot_steps = sum(step_counts.values()) or 1
|
|
self._last_step_mix = {
|
|
"pct_fused": step_counts["fused"] / tot_steps,
|
|
"pct_img": step_counts["image_only"] / tot_steps,
|
|
"pct_md": step_counts["metadata_only"] / tot_steps,
|
|
"phase": "tower_warmup" if in_tower_warmup else ("fused_warmup" if in_fused_warmup else "dynamic")
|
|
}
|
|
|
|
print(
|
|
f"[Epoch {epoch+1}/{self.epochs}] Train Loss: {avg_loss:.4f} | "
|
|
f"Fused Acc: {100*avg_fused:.2f}% | Img Acc: {100*avg_img:.2f}% | Md Acc: {100*avg_md:.2f}%"
|
|
)
|
|
row, stop_now = self.evaluate(epoch)
|
|
if stop_now:
|
|
print(f"[early] stopping on '{self.args.early_metric}' "
|
|
f"with best={self._early.best:.5f} at epoch {epoch+1 - self._early.bad_epochs}")
|
|
break
|
|
|
|
torch.save(self._snapshot(), str(self.models_dir / "model_last.pt"))
|
|
# Ensure holdout-best checkpoint is written
|
|
if getattr(self, "_holdout_best", None):
|
|
hb = self._holdout_best
|
|
if hb.get("state") is not None and hb.get("save_path"):
|
|
try:
|
|
torch.save(hb["state"], hb["save_path"])
|
|
except Exception:
|
|
pass
|
|
|
|
def _copy_best_roc(epoch_idx: int, tag: str):
|
|
if epoch_idx is None or epoch_idx < 0:
|
|
return
|
|
src_dir = self.run_dir / "roc_curves"
|
|
if not src_dir.exists():
|
|
return
|
|
prefix = f"epoch{epoch_idx + 1}_"
|
|
dest_dir = self.run_dir / f"roc_curves_{tag}"
|
|
dest_dir.mkdir(parents=True, exist_ok=True)
|
|
for path in src_dir.glob(f"{prefix}*.json"):
|
|
try:
|
|
shutil.copy2(path, dest_dir / path.name)
|
|
except Exception:
|
|
pass
|
|
if getattr(self, "_early", None):
|
|
print(f"[early] restoring best model at epoch {self._early.best_epoch+1}: {self._early.best:.5f}")
|
|
self._early.restore(self)
|
|
elif getattr(self, "_best_info", None) and self._best_info.get("state") is not None:
|
|
be = self._best_info
|
|
print(f"[best] restoring best model at epoch {be['epoch']+1}: {be['value']:.5f} (monitor={be['monitor']})")
|
|
self._restore_from_state(be["state"])
|
|
torch.save(self._snapshot(), str(self.models_dir / "model_best.pt"))
|
|
# Snapshot ROC curves for both main-best and holdout-best epochs
|
|
if getattr(self, "_early", None) and getattr(self._early, "best_epoch", None) is not None:
|
|
_copy_best_roc(self._early.best_epoch, "best")
|
|
elif getattr(self, "_best_info", None):
|
|
_copy_best_roc(self._best_info.get("epoch", -1), "best")
|
|
if getattr(self, "_holdout_best", None):
|
|
_copy_best_roc(self._holdout_best.get("epoch", -1), "holdout_best")
|
|
|
|
|
|
|
|
def _evaluate_split(self, loader, epoch: int, split: str):
|
|
"""
|
|
Shared evaluation helper that computes core metrics (loss/accuracy/AUC/etc.)
|
|
for a given dataloader. Returns a dict with scalar metrics or None if the
|
|
loader is empty.
|
|
"""
|
|
if loader is None:
|
|
return None
|
|
|
|
total_loss = 0.0
|
|
total_counts = {"fused": 0, "image_only": 0, "metadata_only": 0}
|
|
total_accs = {"fused": 0.0, "image_only": 0.0, "metadata_only": 0.0}
|
|
all_y = []
|
|
probs_f_list, probs_i_list, probs_m_list = [], [], []
|
|
agree_f_img = 0
|
|
agree_f_md = 0
|
|
n_img_comp = 0
|
|
n_md_comp = 0
|
|
|
|
with torch.no_grad():
|
|
for batch in loader:
|
|
imgs, metas, geometry, labels = self._unpack_batch(batch)
|
|
imgs = imgs.to(self.device)
|
|
metas = metas.to(self.device)
|
|
labels = labels.to(self.device)
|
|
geometry = geometry.to(self.device) if geometry is not None else None
|
|
|
|
if self.mode == "vote":
|
|
img_feats = self.img_tower(imgs, geometry)
|
|
md_feats = self.md_tower(metas)
|
|
out_img = self.head_img(img_feats)
|
|
out_md = self.head_md(md_feats)
|
|
out_fused = self.vote(out_img, out_md)
|
|
else:
|
|
img_feats = self.img_tower(imgs, geometry)
|
|
md_feats = self.md_tower(metas)
|
|
outputs = self.bridge(img_feats, md_feats)
|
|
if isinstance(outputs, tuple):
|
|
out_fused, out_img, out_md = outputs
|
|
else:
|
|
out_fused, out_img, out_md = outputs, None, None
|
|
|
|
all_y.append(labels.cpu().numpy())
|
|
probs_f_list.append(F.softmax(out_fused, dim=1).cpu().numpy())
|
|
|
|
loss = self._classification_loss(out_fused, labels)
|
|
total_loss += loss.item()
|
|
|
|
preds_f = out_fused.argmax(dim=1)
|
|
total_accs["fused"] += (preds_f == labels).float().mean().item()
|
|
total_counts["fused"] += 1
|
|
|
|
if out_img is not None:
|
|
preds_i = out_img.argmax(dim=1)
|
|
total_accs["image_only"] += (preds_i == labels).float().mean().item()
|
|
total_counts["image_only"] += 1
|
|
probs_i_list.append(F.softmax(out_img, dim=1).cpu().numpy())
|
|
|
|
agree_f_img += (preds_f == preds_i).sum().item()
|
|
n_img_comp += preds_f.numel()
|
|
|
|
if out_md is not None:
|
|
preds_m = out_md.argmax(dim=1)
|
|
total_accs["metadata_only"] += (preds_m == labels).float().mean().item()
|
|
total_counts["metadata_only"] += 1
|
|
probs_m_list.append(F.softmax(out_md, dim=1).cpu().numpy())
|
|
|
|
agree_f_md += (preds_f == preds_m).sum().item()
|
|
n_md_comp += preds_f.numel()
|
|
|
|
if total_counts["fused"] == 0:
|
|
return None
|
|
|
|
avg_loss = total_loss / total_counts["fused"]
|
|
avg_fused = total_accs["fused"] / total_counts["fused"]
|
|
avg_img = (total_accs["image_only"] / total_counts["image_only"]) if total_counts["image_only"] else 0.0
|
|
avg_md = (total_accs["metadata_only"] / total_counts["metadata_only"]) if total_counts["metadata_only"] else 0.0
|
|
|
|
y_true = np.concatenate(all_y, axis=0) if all_y else np.array([])
|
|
y_prob_f = np.concatenate(probs_f_list, axis=0) if probs_f_list else None
|
|
y_prob_i = np.concatenate(probs_i_list, axis=0) if probs_i_list else None
|
|
y_prob_m = np.concatenate(probs_m_list, axis=0) if probs_m_list else None
|
|
|
|
def compute_multiclass_roc(y_true_np, prob_np):
|
|
if prob_np is None or prob_np.size == 0:
|
|
return None
|
|
C = int(prob_np.shape[1])
|
|
classes = list(range(C))
|
|
y_bin = label_binarize(y_true_np, classes=classes)
|
|
per_class = {}
|
|
aucs = []
|
|
for c in range(C):
|
|
try:
|
|
fpr, tpr, _ = roc_curve(y_bin[:, c], prob_np[:, c])
|
|
auc_val = float(auc(fpr, tpr)) if len(fpr) > 1 else float("nan")
|
|
per_class[c] = {"fpr": fpr.tolist(), "tpr": tpr.tolist(), "auc": auc_val}
|
|
aucs.append(auc_val)
|
|
except Exception:
|
|
per_class[c] = {"fpr": [0.0, 1.0], "tpr": [0.0, 1.0], "auc": float("nan")}
|
|
|
|
try:
|
|
fpr_micro, tpr_micro, _ = roc_curve(y_bin.ravel(), prob_np[:, :C].ravel())
|
|
auc_micro = float(auc(fpr_micro, tpr_micro))
|
|
micro = {"fpr": fpr_micro.tolist(), "tpr": tpr_micro.tolist(), "auc": auc_micro}
|
|
except Exception:
|
|
micro = None
|
|
|
|
auc_vals = [v["auc"] for v in per_class.values() if v["auc"] == v["auc"]]
|
|
macro_auc = float(np.mean(auc_vals)) if auc_vals else float("nan")
|
|
return {"per_class": per_class, "micro": micro, "macro_auc": macro_auc}
|
|
|
|
def macro_ovr_auc(y, p):
|
|
try:
|
|
y = np.asarray(y)
|
|
if p is None:
|
|
return float("nan")
|
|
p = np.asarray(p)
|
|
if p.ndim == 2 and p.shape[1] == 2:
|
|
return roc_auc_score(y, p[:, 1])
|
|
if p.ndim == 1 or p.shape[1] == 1:
|
|
return roc_auc_score(y, p.ravel())
|
|
return roc_auc_score(y, p, multi_class="ovr", average="macro")
|
|
except Exception:
|
|
return float("nan")
|
|
|
|
def top2_acc(y, p):
|
|
if p is None:
|
|
return float("nan")
|
|
if p.ndim != 2:
|
|
return float("nan")
|
|
k = 2 if p.shape[1] >= 2 else 1
|
|
topk = np.argpartition(-p, kth=k - 1, axis=1)[:, :k]
|
|
return np.mean((topk == y[:, None]).any(axis=1).astype(np.float32))
|
|
|
|
def avg_margin(p):
|
|
if p is None or p.ndim != 2:
|
|
return float("nan")
|
|
s = np.sort(p, axis=1)[:, ::-1]
|
|
if s.shape[1] == 1:
|
|
return float("nan")
|
|
return float(np.mean(s[:, 0] - s[:, 1]))
|
|
|
|
roc_f = compute_multiclass_roc(y_true, y_prob_f)
|
|
roc_i = compute_multiclass_roc(y_true, y_prob_i)
|
|
roc_m = compute_multiclass_roc(y_true, y_prob_m)
|
|
|
|
roc_dir = self.run_dir / "roc_curves"
|
|
roc_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
def dump_roc(blob, head: str):
|
|
if blob is None:
|
|
return
|
|
suffix = head if split == "val" else f"{split}_{head}"
|
|
path = roc_dir / f"epoch{epoch + 1}_{suffix}.json"
|
|
with open(path, "w") as f:
|
|
json.dump(blob, f)
|
|
|
|
dump_roc(roc_f, "fused")
|
|
dump_roc(roc_i, "image")
|
|
dump_roc(roc_m, "metadata")
|
|
|
|
auc_f = macro_ovr_auc(y_true, y_prob_f) if y_prob_f is not None else float("nan")
|
|
auc_i = macro_ovr_auc(y_true, y_prob_i) if y_prob_i is not None else float("nan")
|
|
auc_m = macro_ovr_auc(y_true, y_prob_m) if y_prob_m is not None else float("nan")
|
|
|
|
top2_f = top2_acc(y_true, y_prob_f)
|
|
top2_i = top2_acc(y_true, y_prob_i)
|
|
top2_m = top2_acc(y_true, y_prob_m)
|
|
mar_f = avg_margin(y_prob_f)
|
|
mar_i = avg_margin(y_prob_i)
|
|
mar_m = avg_margin(y_prob_m)
|
|
|
|
agree_rate_f_img = (agree_f_img / n_img_comp) if n_img_comp else float("nan")
|
|
agree_rate_f_md = (agree_f_md / n_md_comp) if n_md_comp else float("nan")
|
|
|
|
return {
|
|
"loss": float(avg_loss),
|
|
"acc_fused": float(avg_fused),
|
|
"acc_img": float(avg_img),
|
|
"acc_md": float(avg_md),
|
|
"auc_fused": float(auc_f) if auc_f == auc_f else float("nan"),
|
|
"auc_img": float(auc_i) if auc_i == auc_i else float("nan"),
|
|
"auc_md": float(auc_m) if auc_m == auc_m else float("nan"),
|
|
"top2_fused": float(top2_f) if top2_f == top2_f else float("nan"),
|
|
"top2_img": float(top2_i) if top2_i == top2_i else float("nan"),
|
|
"top2_md": float(top2_m) if top2_m == top2_m else float("nan"),
|
|
"margin_fused": float(mar_f) if mar_f == mar_f else float("nan"),
|
|
"margin_img": float(mar_i) if mar_i == mar_i else float("nan"),
|
|
"margin_md": float(mar_m) if mar_m == mar_m else float("nan"),
|
|
"agree_fused_img": float(agree_rate_f_img) if agree_rate_f_img == agree_rate_f_img else float("nan"),
|
|
"agree_fused_md": float(agree_rate_f_md) if agree_rate_f_md == agree_rate_f_md else float("nan"),
|
|
}
|
|
|
|
|
|
def evaluate(self, epoch: int):
|
|
# Evaluation additionally collects per-head probabilities to report macro AUC
|
|
self.img_tower.eval()
|
|
self.md_tower.eval()
|
|
if self.mode == "vote":
|
|
self.head_img.eval(); self.head_md.eval(); self.vote.eval()
|
|
else:
|
|
self.bridge.eval()
|
|
|
|
metrics_val = self._evaluate_split(self.test_loader, epoch, "val")
|
|
metrics_holdout = self._evaluate_split(self.holdout_loader, epoch, "holdout") if self.holdout_loader else None
|
|
|
|
if metrics_val is None:
|
|
raise RuntimeError("Validation loader produced no batches; cannot compute metrics.")
|
|
|
|
def ffmt(x):
|
|
return "NA" if (x != x) else f"{x:.3f}"
|
|
|
|
print(
|
|
f"[Epoch {epoch+1}/{self.epochs}] Eval Loss: {metrics_val['loss']:.4f} | "
|
|
f"Fused Acc: {100*metrics_val['acc_fused']:.2f}% | "
|
|
f"Img Acc: {100*metrics_val['acc_img']:.2f}% | "
|
|
f"Md Acc: {100*metrics_val['acc_md']:.2f}% | "
|
|
f"Fused AUC: {ffmt(metrics_val['auc_fused'])} | "
|
|
f"Img AUC: {ffmt(metrics_val['auc_img'])} | "
|
|
f"Md AUC: {ffmt(metrics_val['auc_md'])}"
|
|
)
|
|
|
|
if metrics_holdout:
|
|
print(
|
|
f"[Epoch {epoch+1}/{self.epochs}] Holdout Loss: {metrics_holdout['loss']:.4f} | "
|
|
f"Fused Acc: {100*metrics_holdout['acc_fused']:.2f}% | "
|
|
f"Img Acc: {100*metrics_holdout['acc_img']:.2f}% | "
|
|
f"Md Acc: {100*metrics_holdout['acc_md']:.2f}% | "
|
|
f"Fused AUC: {ffmt(metrics_holdout['auc_fused'])} | "
|
|
f"Img AUC: {ffmt(metrics_holdout['auc_img'])} | "
|
|
f"Md AUC: {ffmt(metrics_holdout['auc_md'])}"
|
|
)
|
|
|
|
self.last_eval = {
|
|
"acc_fused": float(metrics_val["acc_fused"]),
|
|
"acc_img": float(metrics_val["acc_img"]),
|
|
"acc_md": float(metrics_val["acc_md"]),
|
|
"auc_fused": float(metrics_val["auc_fused"]) if metrics_val["auc_fused"] == metrics_val["auc_fused"] else 0.0,
|
|
"auc_img": float(metrics_val["auc_img"]) if metrics_val["auc_img"] == metrics_val["auc_img"] else 0.0,
|
|
"auc_md": float(metrics_val["auc_md"]) if metrics_val["auc_md"] == metrics_val["auc_md"] else 0.0,
|
|
}
|
|
|
|
row = {
|
|
"epoch": int(epoch + 1),
|
|
"eval_loss": _num(metrics_val["loss"]),
|
|
"acc_fused": _num(metrics_val["acc_fused"]),
|
|
"acc_img": _num(metrics_val["acc_img"]),
|
|
"acc_md": _num(metrics_val["acc_md"]),
|
|
"auc_fused": _num(metrics_val["auc_fused"]),
|
|
"auc_img": _num(metrics_val["auc_img"]),
|
|
"auc_md": _num(metrics_val["auc_md"]),
|
|
"top2_fused": _num(metrics_val["top2_fused"]),
|
|
"top2_img": _num(metrics_val["top2_img"]),
|
|
"top2_md": _num(metrics_val["top2_md"]),
|
|
"margin_fused": _num(metrics_val["margin_fused"]),
|
|
"margin_img": _num(metrics_val["margin_img"]),
|
|
"margin_md": _num(metrics_val["margin_md"]),
|
|
"agree_fused_img": _num(metrics_val["agree_fused_img"]),
|
|
"agree_fused_md": _num(metrics_val["agree_fused_md"]),
|
|
}
|
|
|
|
if metrics_holdout:
|
|
row.update({
|
|
"holdout_loss": _num(metrics_holdout["loss"]),
|
|
"holdout_acc_fused": _num(metrics_holdout["acc_fused"]),
|
|
"holdout_acc_img": _num(metrics_holdout["acc_img"]),
|
|
"holdout_acc_md": _num(metrics_holdout["acc_md"]),
|
|
"holdout_auc_fused": _num(metrics_holdout["auc_fused"]),
|
|
"holdout_auc_img": _num(metrics_holdout["auc_img"]),
|
|
"holdout_auc_md": _num(metrics_holdout["auc_md"]),
|
|
})
|
|
else:
|
|
row.setdefault("holdout_loss", None)
|
|
row.setdefault("holdout_acc_fused", None)
|
|
row.setdefault("holdout_acc_img", None)
|
|
row.setdefault("holdout_acc_md", None)
|
|
row.setdefault("holdout_auc_fused", None)
|
|
row.setdefault("holdout_auc_img", None)
|
|
row.setdefault("holdout_auc_md", None)
|
|
|
|
# realized step mix from training (if present)
|
|
mix = getattr(self, "_last_step_mix", None)
|
|
if mix:
|
|
row.update({
|
|
"pct_fused": _num(mix.get("pct_fused")),
|
|
"pct_img": _num(mix.get("pct_img")),
|
|
"pct_md": _num(mix.get("pct_md")),
|
|
"phase": mix.get("phase"),
|
|
})
|
|
|
|
# SE gate stats (if the bridge exposes them)
|
|
# fetch without clearing
|
|
bridge_module = getattr(self, "bridge", None)
|
|
get_se_stats = getattr(bridge_module, "get_se_stats", None) if bridge_module is not None else None
|
|
if callable(get_se_stats):
|
|
se_stats = get_se_stats(reset=False)
|
|
if se_stats:
|
|
row.update({
|
|
"se_mean": _num(se_stats.get("mean")),
|
|
"se_std": _num(se_stats.get("std")),
|
|
"se_pct_lt_0.2": _num(se_stats.get("pct_lt_0.2")),
|
|
"se_pct_gt_0.8": _num(se_stats.get("pct_gt_0.8")),
|
|
})
|
|
# debug print BEFORE reset so you can see the true count
|
|
try:
|
|
print(f"SE gate samples seen: {getattr(getattr(bridge_module, 'se_log', None), '_n', 0)}")
|
|
except Exception:
|
|
pass
|
|
# now clear for the next epoch
|
|
try:
|
|
get_se_stats(reset=True)
|
|
except Exception:
|
|
pass
|
|
row.setdefault(self.args.early_metric, None)
|
|
# early-stop decision (no break here)
|
|
should_stop = False
|
|
if getattr(self, "_early", None):
|
|
should_stop = self._early.step(row, trainer=self, epoch=epoch)
|
|
row.update({
|
|
"early_best_so_far" : float(self._early.best),
|
|
"early_bad_epochs" : int(self._early.bad_epochs),
|
|
"early_improved" : int(self._early.last_improved),
|
|
"early_monitor" : self.args.early_metric,
|
|
})
|
|
# concise console status each epoch
|
|
mon = self.args.early_metric
|
|
cur = row.get(mon, None)
|
|
if cur is not None and cur == cur: # not NaN
|
|
msg = (
|
|
f"[early] epoch {epoch+1}: {mon}={cur:.5f} | best={self._early.best:.5f} | "
|
|
f"bad_epochs={self._early.bad_epochs}/{self._early.patience} | improved={'yes' if self._early.last_improved else 'no'}"
|
|
)
|
|
print(msg)
|
|
try:
|
|
self.logger.info(msg)
|
|
except Exception:
|
|
pass
|
|
|
|
# Holdout best-tracker (if a holdout set is present)
|
|
if getattr(self, "_holdout_best", None):
|
|
mon_h = self._holdout_best["monitor"]
|
|
mode_h = self._holdout_best["mode"]
|
|
min_delta_h = self._holdout_best["min_delta"]
|
|
val_h = row.get(mon_h, None)
|
|
improved_h = False
|
|
if val_h is not None and val_h == val_h:
|
|
best_val_h = self._holdout_best["value"]
|
|
if mode_h == "max":
|
|
improved_h = (val_h > best_val_h + min_delta_h)
|
|
else:
|
|
improved_h = (val_h < best_val_h - min_delta_h)
|
|
if improved_h:
|
|
self._holdout_best["value"] = float(val_h)
|
|
self._holdout_best["epoch"] = int(epoch)
|
|
st_h = self._snapshot()
|
|
self._holdout_best["state"] = st_h
|
|
if self._holdout_best.get("save_path"):
|
|
try:
|
|
torch.save(st_h, self._holdout_best["save_path"])
|
|
except Exception:
|
|
pass
|
|
print(f"[holdout_best] ↑ new best {mon_h}={val_h:.5f} at epoch {epoch+1}")
|
|
row.update({
|
|
"holdout_best_monitor": mon_h,
|
|
"holdout_best_so_far": (None if self._holdout_best["value"] in (math.inf, -math.inf) else float(self._holdout_best["value"])),
|
|
"holdout_best_epoch": (None if self._holdout_best["epoch"] < 0 else int(self._holdout_best["epoch"] + 1)),
|
|
})
|
|
|
|
# Best-tracker (always-on): save best model snapshot and optional checkpoint
|
|
if getattr(self, "_best_info", None):
|
|
# If early stopper is active, mirror its best info and avoid duplicate prints
|
|
if getattr(self, "_early", None):
|
|
row.update({
|
|
"best_monitor": self._best_info["monitor"],
|
|
"best_so_far": float(self._early.best) if self._early.best == self._early.best else None,
|
|
"best_epoch": (int(self._early.best_epoch) + 1) if (self._early.best_epoch is not None) else None,
|
|
})
|
|
# skip independent tracking/prints to avoid duplication
|
|
self._write_epoch_log(row)
|
|
return row, should_stop
|
|
mon = self._best_info["monitor"]
|
|
mode = self._best_info["mode"]
|
|
min_delta = self._best_info["min_delta"]
|
|
val = row.get(mon, None)
|
|
improved = False
|
|
if val is not None and val == val: # not NaN
|
|
best_val = self._best_info["value"]
|
|
if mode == "max":
|
|
improved = (val > best_val + min_delta)
|
|
else:
|
|
improved = (val < best_val - min_delta)
|
|
if improved:
|
|
self._best_info["value"] = float(val)
|
|
self._best_info["epoch"] = int(epoch)
|
|
st = self._snapshot()
|
|
self._best_info["state"] = st
|
|
if self._best_info.get("save_path"):
|
|
try:
|
|
torch.save(st, self._best_info["save_path"])
|
|
except Exception:
|
|
pass
|
|
print(f"[best] ↑ new best {mon}={val:.5f} at epoch {epoch+1}")
|
|
# Add best-so-far to row
|
|
row.update({
|
|
"best_monitor": mon,
|
|
"best_so_far": (None if self._best_info["value"] in (math.inf, -math.inf) else float(self._best_info["value"])),
|
|
"best_epoch": (None if self._best_info["epoch"] < 0 else int(self._best_info["epoch"] + 1)),
|
|
})
|
|
|
|
# finally: write once, after all updates
|
|
self._write_epoch_log(row)
|
|
return row, should_stop
|
|
|
|
def _restore_from_state(self, st: dict):
|
|
if not st:
|
|
return
|
|
self.img_tower.load_state_dict(st["img_tower"])
|
|
self.md_tower.load_state_dict(st["md_tower"])
|
|
if "bridge" in st and hasattr(self, "bridge"):
|
|
self.bridge.load_state_dict(st["bridge"])
|
|
if "head_img" in st and hasattr(self, "head_img"):
|
|
self.head_img.load_state_dict(st["head_img"])
|
|
if "head_md" in st and hasattr(self, "head_md"):
|
|
self.head_md.load_state_dict(st["head_md"])
|
|
if "optimizer" in st:
|
|
self.optimizer.load_state_dict(st["optimizer"])
|