from __future__ import annotations import csv import json import time from copy import deepcopy from datetime import datetime, timezone from pathlib import Path from typing import Any import torch from torch.utils.data import DataLoader from tqdm import tqdm from datasets.cd_dataset import CDDataset from utils.config_loader import load_dataset_config, load_model_config from utils.dataset_cache import apply_dataloader_cli_overrides, dataloader_kwargs, dataloader_policy_lines, dataset_runtime_summary, print_dataloader_policy from utils.gpu_utils import print_gpu_diagnostics, resolve_gpu from utils.metrics import BinaryMetrics, BoundaryMetrics, normalize_binary_prediction from utils.model_adapters import BaseModelAdapter, get_model_adapter from utils.profiling import GpuProfiler, ProfilingUnavailable, count_flops, count_parameters from utils.qualitative import ( denormalize, manifest_ids, rank_for_sample, safe_sample_id, save_binary_prediction, save_probability_map, save_visual_panel, select_or_load_manifest, ) from utils.results_writer import append_to_comparison_table, save_metrics ROOT = Path(__file__).resolve().parents[1] def _json_dump(path: Path, payload: Any) -> None: path.parent.mkdir(parents=True, exist_ok=True) with path.open("w", encoding="utf-8") as f: json.dump(payload, f, indent=2, sort_keys=True) def _append_jsonl(path: Path, payload: dict) -> None: path.parent.mkdir(parents=True, exist_ok=True) with path.open("a", encoding="utf-8") as f: f.write(json.dumps(payload, sort_keys=True) + "\n") def _threshold_sweep_enabled(dataset_cfg: dict, split: str) -> bool: eval_cfg = dataset_cfg.get("eval", {}) if split == "val": return bool(eval_cfg.get("sweep_val_threshold", True)) if split == "test": return bool(eval_cfg.get("sweep_test_threshold", False)) return bool(eval_cfg.get("sweep_threshold", True)) def _thresholds(dataset_cfg: dict, split: str) -> list[float]: eval_cfg = dataset_cfg.get("eval", {}) if not _threshold_sweep_enabled(dataset_cfg, split): return [float(eval_cfg.get("threshold", 0.5))] if "thresholds" in eval_cfg: return [float(x) for x in eval_cfg["thresholds"]] start = float(eval_cfg.get("threshold_min", 0.05)) stop = float(eval_cfg.get("threshold_max", 0.95)) step = float(eval_cfg.get("threshold_step", 0.05)) values = [] current = start while current <= stop + 1e-9: values.append(round(current, 4)) current += step return values def _cfg_with_threshold(dataset_cfg: dict, threshold: float) -> dict: cfg = deepcopy(dataset_cfg) cfg.setdefault("eval", {})["threshold"] = float(threshold) return cfg def _last_tensor(raw_output: Any) -> torch.Tensor | None: value = raw_output[-1] if isinstance(raw_output, (list, tuple)) and raw_output else raw_output return value if torch.is_tensor(value) else None def _has_nonfinite_tensor(value: Any) -> bool: if torch.is_tensor(value): return not bool(torch.isfinite(value.detach()).all().item()) if isinstance(value, (list, tuple)): return any(_has_nonfinite_tensor(item) for item in value) if isinstance(value, dict): return any(_has_nonfinite_tensor(item) for item in value.values()) return False def _batch_stats(raw_output: Any, normalized, mask: torch.Tensor) -> dict[str, float | None]: binary = normalized.binary.float() target = (mask > 0).float() score = normalized.score raw_tensor = _last_tensor(raw_output) return { "gt_positive_ratio": float(target.mean().item()), "pred_positive_ratio": float(binary.mean().item()), "mean_prob": float(score.mean().item()) if score is not None else None, "logit_mean": float(raw_tensor.detach().float().mean().item()) if raw_tensor is not None else None, "logit_std": float(raw_tensor.detach().float().std().item()) if raw_tensor is not None and raw_tensor.numel() > 1 else None, } def _loader(dataset_cfg: dict, split: str, batch_size: int, shuffle: bool, drop_last: bool = False) -> DataLoader: ds = CDDataset(dataset_cfg["data_root"], split, cfg=dataset_cfg, return_format="tuple") return DataLoader( ds, batch_size=batch_size, shuffle=shuffle, **dataloader_kwargs(dataset_cfg, torch.cuda.is_available()), drop_last=drop_last, ) def _write_trajectory(metrics_dir: Path, trajectory: list[dict]) -> None: _json_dump(metrics_dir / "trajectory.json", trajectory) if not trajectory: return columns = [ "epoch", "train_loss", "val_threshold", "val_f1", "val_iou", "val_miou", "val_precision", "val_recall", "val_oa", "val_bf1", "val_gt_positive_ratio", "val_pred_positive_ratio", "val_mean_prob", "diagnostic_test_f1", "diagnostic_test_iou", "val_test_f1_gap", "val_test_iou_gap", "best_val_f1_so_far", "best_val_epoch_or_iter_so_far", "checkpoint_path", ] with (metrics_dir / "trajectory.csv").open("w", newline="", encoding="utf-8") as f: writer = csv.DictWriter(f, fieldnames=columns, extrasaction="ignore") writer.writeheader() writer.writerows(trajectory) def _evaluate_split( *, model_name: str, model: torch.nn.Module, adapter: BaseModelAdapter, dataset_cfg: dict, loader: DataLoader, device: torch.device, threshold: float, split: str, out_dir: Path, save_outputs: bool, checkpoint_path: Path | None = None, ) -> dict: cfg = _cfg_with_threshold(dataset_cfg, threshold) boundary = BoundaryMetrics(tolerance=int(dataset_cfg.get("eval", {}).get("boundary_tolerance", 2))) metrics = BinaryMetrics(threshold=threshold) pred_dir = out_dir / "predictions" / split prob_dir = out_dir / "predictions" / f"{split}_prob" visual_dir = out_dir / "visuals" / "selected_20" manifest = select_or_load_manifest(dataset_cfg) if save_outputs and split == "test" else {"samples": []} selected = manifest_ids(manifest) mean_a = dataset_cfg.get("mean_a", [0.485, 0.456, 0.406]) std_a = dataset_cfg.get("std_a", [0.229, 0.224, 0.225]) mean_b = dataset_cfg.get("mean_b", mean_a) std_b = dataset_cfg.get("std_b", std_a) n_samples = 0 model_time = 0.0 end_to_end_start = time.perf_counter() stat_sums = {"gt_positive_ratio": 0.0, "pred_positive_ratio": 0.0, "mean_prob": 0.0} stat_counts = {"mean_prob": 0} model.eval() with torch.inference_mode(), GpuProfiler(device=device, required=False) as gpu_profiler: for batch in loader: if device.type == "cuda": torch.cuda.synchronize(device) start = time.perf_counter() raw = adapter.forward(model, batch, device) if device.type == "cuda": torch.cuda.synchronize(device) elapsed = time.perf_counter() - start a, b, mask, names = batch normalized = adapter.normalize_output(raw, batch, cfg) metrics.update(normalized.metric_tensor, mask) boundary.update(normalized.binary, mask) n_samples += int(mask.shape[0]) model_time += elapsed stats = _batch_stats(raw, normalized, mask) stat_sums["gt_positive_ratio"] += float(stats["gt_positive_ratio"] or 0.0) * int(mask.shape[0]) stat_sums["pred_positive_ratio"] += float(stats["pred_positive_ratio"] or 0.0) * int(mask.shape[0]) if stats["mean_prob"] is not None: stat_sums["mean_prob"] += float(stats["mean_prob"]) * int(mask.shape[0]) stat_counts["mean_prob"] += int(mask.shape[0]) if save_outputs: for i, sample_id in enumerate(names): clean_id = safe_sample_id(str(sample_id)) pred_i = normalized.binary[i] save_binary_prediction(pred_i, pred_dir / f"{clean_id}_pred.png") prob_i = normalized.score[i] if normalized.score is not None else None if prob_i is not None: save_probability_map(prob_i, prob_dir / f"{clean_id}_prob.png") if split == "test" and str(sample_id) in selected: rank = rank_for_sample(manifest, str(sample_id)) save_visual_panel( denormalize(a[i].detach().cpu(), mean_a, std_a), denormalize(b[i].detach().cpu(), mean_b, std_b), mask[i], pred_i, visual_dir / f"{rank:02d}_{clean_id}_panel.png", prob=prob_i, ) elapsed_total = time.perf_counter() - end_to_end_start result = metrics.compute() result.update(boundary.compute()) result.update(gpu_profiler.summary()) result.update({ "model": model_name, "dataset": dataset_cfg["name"], "split": split, "threshold": threshold, "threshold_mode": adapter.get_threshold_mode(), "gt_positive_ratio": stat_sums["gt_positive_ratio"] / max(n_samples, 1), "pred_positive_ratio": stat_sums["pred_positive_ratio"] / max(n_samples, 1), "mean_prob": stat_sums["mean_prob"] / stat_counts["mean_prob"] if stat_counts["mean_prob"] else None, "fps": n_samples / model_time if model_time > 0 else None, "fps_model_only": n_samples / model_time if model_time > 0 else None, "fps_end_to_end": n_samples / elapsed_total if elapsed_total > 0 else None, "sample_count": n_samples, "test_sample_count": n_samples if split == "test" else None, "checkpoint": str(checkpoint_path) if checkpoint_path else None, "timestamp": datetime.now(timezone.utc).isoformat(), "status": "complete", }) return result def _metrics_from_cached_predictions( *, model_name: str, dataset_cfg: dict, cached: list[tuple[torch.Tensor, torch.Tensor]], threshold: float, split: str, n_samples: int, model_time: float, elapsed_total: float, gpu_summary: dict[str, object], ) -> dict: metrics = BinaryMetrics(threshold=threshold) boundary = BoundaryMetrics(tolerance=int(dataset_cfg.get("eval", {}).get("boundary_tolerance", 2))) stat_sums = {"gt_positive_ratio": 0.0, "pred_positive_ratio": 0.0, "mean_prob": 0.0} stat_counts = {"mean_prob": 0} for metric_tensor, mask in cached: binary, score = normalize_binary_prediction(metric_tensor, threshold=threshold) metrics.update(metric_tensor, mask) boundary.update(binary, mask) batch_size = int(mask.shape[0]) target = (mask > 0).float() stat_sums["gt_positive_ratio"] += float(target.mean().item()) * batch_size stat_sums["pred_positive_ratio"] += float(binary.float().mean().item()) * batch_size if score is not None: stat_sums["mean_prob"] += float(score.float().mean().item()) * batch_size stat_counts["mean_prob"] += batch_size result = metrics.compute() result.update(boundary.compute()) result.update(gpu_summary) result.update({ "model": model_name, "dataset": dataset_cfg["name"], "split": split, "threshold": threshold, "threshold_mode": "threshold", "gt_positive_ratio": stat_sums["gt_positive_ratio"] / max(n_samples, 1), "pred_positive_ratio": stat_sums["pred_positive_ratio"] / max(n_samples, 1), "mean_prob": stat_sums["mean_prob"] / stat_counts["mean_prob"] if stat_counts["mean_prob"] else None, "fps": n_samples / model_time if model_time > 0 else None, "fps_model_only": n_samples / model_time if model_time > 0 else None, "fps_end_to_end": n_samples / elapsed_total if elapsed_total > 0 else None, "sample_count": n_samples, "test_sample_count": n_samples if split == "test" else None, "checkpoint": None, "timestamp": datetime.now(timezone.utc).isoformat(), "status": "complete", "sweep_cache": True, }) return result def _threshold_sweep( *, model_name: str, model: torch.nn.Module, adapter: BaseModelAdapter, dataset_cfg: dict, loader: DataLoader, device: torch.device, split: str, out_dir: Path, ) -> dict: if adapter.get_threshold_mode() == "argmax": threshold = adapter.get_threshold(dataset_cfg) metrics = _evaluate_split( model_name=model_name, model=model, adapter=adapter, dataset_cfg=dataset_cfg, loader=loader, device=device, threshold=threshold, split=split, out_dir=out_dir, save_outputs=False, ) return { "threshold_mode": "argmax", "selected_threshold": threshold, "best_f1": metrics["f1"], "results": [metrics], } cached: list[tuple[torch.Tensor, torch.Tensor]] = [] n_samples = 0 model_time = 0.0 end_to_end_start = time.perf_counter() model.eval() with torch.inference_mode(), GpuProfiler(device=device, required=False) as gpu_profiler: for batch in loader: if device.type == "cuda": torch.cuda.synchronize(device) start = time.perf_counter() raw = adapter.forward(model, batch, device) if device.type == "cuda": torch.cuda.synchronize(device) model_time += time.perf_counter() - start _a, _b, mask, _names = batch normalized = adapter.normalize_output(raw, batch, dataset_cfg) cached.append((normalized.metric_tensor.detach().cpu(), mask.detach().cpu())) n_samples += int(mask.shape[0]) elapsed_total = time.perf_counter() - end_to_end_start gpu_summary = gpu_profiler.summary() results = [] best = None for threshold in _thresholds(dataset_cfg, split): metrics = _metrics_from_cached_predictions( model_name=model_name, dataset_cfg=dataset_cfg, threshold=threshold, split=split, cached=cached, n_samples=n_samples, model_time=model_time, elapsed_total=elapsed_total, gpu_summary=gpu_summary, ) results.append(metrics) if best is None or metrics["f1"] > best["f1"]: best = metrics assert best is not None return { "threshold_mode": "threshold", "selected_threshold": float(best["threshold"]), "best_f1": best["f1"], "results": results, } def train_with_adapter(model_name: str, args) -> int: dataset_cfg = load_dataset_config(args.dataset) apply_dataloader_cli_overrides(dataset_cfg, args) model_cfg = load_model_config(model_name) if args.epochs is not None: model_cfg["num_epochs"] = int(args.epochs) if args.lr is not None: model_cfg["lr"] = float(args.lr) batch_size = int(args.batch_size or dataset_cfg.get("batch_size", 8)) adapter = get_model_adapter(model_name) if not adapter.supports_unified_training: raise RuntimeError(f"{model_name} does not support unified training: {adapter.notes_or_failure_reason}") if args.dry_run: print(f"[DRY-RUN] unified model={model_name} dataset={dataset_cfg['name']} root={dataset_cfg['data_root']}") print(f"[DATASET] {dataset_runtime_summary(dataset_cfg)}") print_dataloader_policy(dataset_cfg, torch.cuda.is_available()) return 0 gpu_resolution = resolve_gpu(args.gpu) print_gpu_diagnostics(gpu_resolution) device = torch.device(gpu_resolution.local_device) out_dir = Path(args.output_dir) if args.output_dir else ROOT / "results" / model_name / dataset_cfg["name"] if not out_dir.is_absolute(): out_dir = ROOT / out_dir ckpt_dir = out_dir / "checkpoints" log_dir = out_dir / "logs" metrics_dir = out_dir / "metrics" for path in (ckpt_dir, log_dir, metrics_dir, out_dir / "predictions" / "test", out_dir / "visuals" / "selected_20"): path.mkdir(parents=True, exist_ok=True) train_loader = _loader(dataset_cfg, "train", batch_size, shuffle=True, drop_last=True) val_loader = _loader(dataset_cfg, "val", batch_size, shuffle=False) test_loader = _loader(dataset_cfg, "test", batch_size, shuffle=False) print_dataloader_policy(dataset_cfg, torch.cuda.is_available()) model = adapter.build_model(model_cfg, dataset_cfg, device) optimizer = adapter.build_optimizer(model, model_cfg) scheduler = adapter.build_scheduler(optimizer, model_cfg) amp_enabled = device.type == "cuda" and bool(getattr(adapter, "supports_amp_training", True)) scaler = torch.amp.GradScaler("cuda", enabled=amp_enabled) epochs = int(model_cfg.get("num_epochs", 200)) best_path = ckpt_dir / "best_model.pth" latest_path = ckpt_dir / "latest.pth" best_f1 = -1.0 best_epoch = 0 selected_threshold = adapter.get_threshold(dataset_cfg) trajectory: list[dict] = [] train_history: list[dict] = [] val_history: list[dict] = [] start_epoch = 1 if args.resume and latest_path.exists(): checkpoint = torch.load(latest_path, map_location=device) adapter.load_checkpoint(model, latest_path, device) if isinstance(checkpoint, dict) and "optimizer_state_dict" in checkpoint: optimizer.load_state_dict(checkpoint["optimizer_state_dict"]) if scheduler is not None and isinstance(checkpoint, dict) and "scheduler_state_dict" in checkpoint: scheduler.load_state_dict(checkpoint["scheduler_state_dict"]) start_epoch = int(checkpoint.get("epoch", 0)) + 1 best_f1 = float(checkpoint.get("best_val_f1", best_f1)) best_epoch = int(checkpoint.get("best_epoch", 0)) selected_threshold = float(checkpoint.get("selected_threshold", selected_threshold)) if args.eval_only: adapter.load_checkpoint(model, best_path, device) test_metrics = _evaluate_split( model_name=model_name, model=model, adapter=adapter, dataset_cfg=dataset_cfg, loader=test_loader, device=device, threshold=selected_threshold, split="test", out_dir=out_dir, save_outputs=True, checkpoint_path=best_path, ) save_metrics(model_name, dataset_cfg["name"], "test", test_metrics) _json_dump(metrics_dir / "test_metrics.json", test_metrics) append_to_comparison_table() return 0 train_log = log_dir / "train.log" with train_log.open("w", encoding="utf-8") as log: log.write(f"# Unified Mamba-CD-style training: {model_name}/{dataset_cfg['name']}\n") log.write(f"# Dataset: {dataset_runtime_summary(dataset_cfg)}\n") for line in dataloader_policy_lines(dataset_cfg, torch.cuda.is_available()): log.write(line + "\n") if dataset_cfg.get("io_warning"): print(f"[DATASET-WARNING] {dataset_cfg['io_warning']}") log.write(f"# Warning: {dataset_cfg['io_warning']}\n") for epoch in range(start_epoch, epochs + 1): model.train() total_loss = 0.0 component_sums: dict[str, float] = {} batch_count = 0 progress = tqdm(train_loader, desc=f"{model_name} epoch {epoch}/{epochs}", dynamic_ncols=True) for batch_idx, batch in enumerate(progress, start=1): optimizer.zero_grad(set_to_none=True) with torch.amp.autocast("cuda", enabled=amp_enabled): raw = adapter.forward(model, batch, device) if _has_nonfinite_tensor(raw): raise RuntimeError( f"{model_name}/{dataset_cfg['name']} produced non-finite model output " f"at epoch {epoch}, batch {batch_idx}." ) with torch.amp.autocast("cuda", enabled=False): loss_dict = adapter.compute_loss(raw, batch, model_cfg, dataset_cfg, device) loss = loss_dict["loss"] if not bool(torch.isfinite(loss.detach()).all().item()): components = { key: float(value.detach().float().cpu().item()) for key, value in loss_dict.items() if torch.is_tensor(value) and value.numel() == 1 } raise RuntimeError( f"{model_name}/{dataset_cfg['name']} produced non-finite loss " f"at epoch {epoch}, batch {batch_idx}: {components}" ) scaler.scale(loss).backward() clip = model_cfg.get("grad_clip", model_cfg.get("gradient_clip", None)) if clip is not None: scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), float(clip)) scaler.step(optimizer) scaler.update() total_loss += float(loss.detach().cpu().item()) batch_count += 1 for key, value in loss_dict.items(): component_sums[key] = component_sums.get(key, 0.0) + float(value.detach().cpu().item()) lr = optimizer.param_groups[0]["lr"] gpu_mem = torch.cuda.max_memory_allocated(device) / (1024.0 ** 3) if device.type == "cuda" else 0.0 progress.set_postfix(loss=f"{loss.item():.4f}", lr=f"{lr:.2e}", best_f1=f"{best_f1:.4f}", gpu_gb=f"{gpu_mem:.2f}") if scheduler is not None: scheduler.step() train_loss = total_loss / max(batch_count, 1) train_row = { "epoch": epoch, "train_loss": train_loss, "lr": optimizer.param_groups[0]["lr"], "timestamp": datetime.now(timezone.utc).isoformat(), } for key, value in component_sums.items(): train_row[key] = value / max(batch_count, 1) train_history.append(train_row) sweep = _threshold_sweep( model_name=model_name, model=model, adapter=adapter, dataset_cfg=dataset_cfg, loader=val_loader, device=device, split="val", out_dir=out_dir, ) selected_threshold = float(sweep["selected_threshold"]) _json_dump(metrics_dir / f"val_threshold_sweep_epoch_{epoch}.json", sweep) val_metrics = dict(max(sweep["results"], key=lambda row: row.get("f1", -1.0))) val_metrics["checkpoint"] = str(latest_path) is_best = val_metrics["f1"] > best_f1 if is_best: best_f1 = float(val_metrics["f1"]) best_epoch = epoch metadata = { "model": model_name, "dataset": dataset_cfg["name"], "epoch": epoch, "best_epoch": best_epoch, "best_val_f1": best_f1, "selected_threshold": selected_threshold, "threshold_mode": adapter.get_threshold_mode(), "timestamp": datetime.now(timezone.utc).isoformat(), } adapter.save_checkpoint(model, optimizer, scheduler, latest_path, metadata) if is_best: adapter.save_checkpoint(model, optimizer, scheduler, best_path, metadata) val_payload = dict(val_metrics) val_payload.update({"epoch": epoch, "best_epoch": best_epoch, "selected_threshold": selected_threshold}) val_history.append(val_payload) save_metrics(model_name, dataset_cfg["name"], "val", val_payload) _json_dump(metrics_dir / "val_history.json", val_history) _json_dump(metrics_dir / "train_history.json", train_history) trajectory_row = { "epoch": epoch, "train_loss": train_loss, "val_threshold": selected_threshold, "val_f1": val_payload["f1"], "val_iou": val_payload["iou"], "val_miou": val_payload["miou"], "val_precision": val_payload["precision"], "val_recall": val_payload["recall"], "val_oa": val_payload["oa"], "val_bf1": val_payload.get("bf1"), "val_gt_positive_ratio": val_payload.get("gt_positive_ratio"), "val_pred_positive_ratio": val_payload.get("pred_positive_ratio"), "val_mean_prob": val_payload.get("mean_prob"), "diagnostic_test_f1": None, "diagnostic_test_iou": None, "val_test_f1_gap": None, "val_test_iou_gap": None, "best_val_f1_so_far": best_f1, "best_val_epoch_or_iter_so_far": best_epoch, "checkpoint_path": str(best_path if is_best else latest_path), } trajectory.append(trajectory_row) _write_trajectory(metrics_dir, trajectory) warning = "" if val_payload["recall"] == 0: warning += " zero_recall" if val_payload.get("pred_positive_ratio", 0) == 0: warning += " all_background_prediction" line = ( f"[Epoch {epoch:03d}/{epochs:03d}] train_loss={train_loss:.4f} " f"val_F1={val_payload['f1']:.4f} val_IoU={val_payload['iou']:.4f} " f"val_Prec={val_payload['precision']:.4f} val_Rec={val_payload['recall']:.4f} " f"threshold={selected_threshold} best_F1={best_f1:.4f}{' BEST' if is_best else ''}{warning}" ) print(line, flush=True) log.write(line + "\n") log.flush() adapter.load_checkpoint(model, best_path, device) checkpoint = torch.load(best_path, map_location=device) selected_threshold = float(checkpoint.get("selected_threshold", selected_threshold)) if isinstance(checkpoint, dict) else selected_threshold test_metrics = _evaluate_split( model_name=model_name, model=model, adapter=adapter, dataset_cfg=dataset_cfg, loader=test_loader, device=device, threshold=selected_threshold, split="test", out_dir=out_dir, save_outputs=True, checkpoint_path=best_path, ) test_metrics.update(count_parameters(model)) try: if not adapter.supports_flops: raise ProfilingUnavailable(f"{model_name} adapter does not support FLOPs.") test_metrics.update(count_flops(model, lambda: adapter.get_dummy_inputs(dataset_cfg, device), device)) except ProfilingUnavailable as exc: test_metrics.update({"flops": None, "flops_g": None, "flops_error": str(exc)}) test_metrics.update({ "best_epoch": best_epoch, "selected_threshold": selected_threshold, "checkpoint_path": str(best_path), "official_selection_metric": "validation_f1", }) save_metrics(model_name, dataset_cfg["name"], "test", test_metrics) _json_dump(metrics_dir / "test_metrics.json", test_metrics) if _threshold_sweep_enabled(dataset_cfg, "test"): test_sweep = _threshold_sweep( model_name=model_name, model=model, adapter=adapter, dataset_cfg=dataset_cfg, loader=test_loader, device=device, split="test", out_dir=out_dir, ) test_sweep.update({ "diagnostic_only": True, "warning": "Test threshold sweep is diagnostic only and was not used for model selection.", "official_validation_selected_threshold": selected_threshold, "official_test_metrics": test_metrics, }) else: test_sweep = { "diagnostic_only": True, "skipped": True, "reason": "Set eval.sweep_test_threshold: true in the dataset config to enable diagnostic test threshold sweeps.", "official_validation_selected_threshold": selected_threshold, "official_test_metrics": test_metrics, } _json_dump(metrics_dir / "test_threshold_sweep.json", test_sweep) best_val = max(val_history, key=lambda row: row.get("f1", -1.0)) if val_history else {} overfit = { "best_validation_epoch": best_epoch, "best_validation_f1": best_f1, "official_test_f1_at_validation_best_checkpoint": test_metrics.get("f1"), "validation_test_f1_gap": (best_val.get("f1", 0.0) - test_metrics.get("f1", 0.0)) if best_val else None, "validation_test_iou_gap": (best_val.get("iou", 0.0) - test_metrics.get("iou", 0.0)) if best_val else None, "gap_warning": bool(best_val and (best_val.get("f1", 0.0) - test_metrics.get("f1", 0.0)) > 0.1), "diagnostic_test_during_training_enabled": False, "note": "Test metrics were not used for checkpoint selection.", } _json_dump(metrics_dir / "overfit_diagnostics.json", overfit) if trajectory: trajectory[-1]["val_test_f1_gap"] = overfit["validation_test_f1_gap"] trajectory[-1]["val_test_iou_gap"] = overfit["validation_test_iou_gap"] _write_trajectory(metrics_dir, trajectory) append_to_comparison_table() _append_jsonl(ROOT / "results" / "training_log.jsonl", { "model": model_name, "dataset": dataset_cfg["name"], "status": "complete", "best_epoch": best_epoch, "best_val_f1": best_f1, "test_f1": test_metrics.get("f1"), "timestamp": datetime.now(timezone.utc).isoformat(), "trainer": "unified", }) return 0