From a97b4da1ad8782c597723b243bd98011d58dd03d Mon Sep 17 00:00:00 2001 From: CaoWangrenbo Date: Sat, 1 Aug 2026 17:47:39 +0800 Subject: [PATCH] fix: persist and restore event threshold via checkpoint Save event_threshold to checkpoint dict during training, restore it during resume and evaluation. evaluate.py now reads from checkpoint instead of hardcoding train_cfg default, so evaluation matches the threshold used during training. Generated by Mistral Vibe. Co-Authored-By: Mistral Vibe --- src/velocity_prediction/evaluate.py | 6 +++++- src/velocity_prediction/train.py | 6 ++++++ 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/src/velocity_prediction/evaluate.py b/src/velocity_prediction/evaluate.py index 7aa96c1..4fe33c4 100644 --- a/src/velocity_prediction/evaluate.py +++ b/src/velocity_prediction/evaluate.py @@ -166,7 +166,11 @@ def main(): ckpt = torch.load(args.checkpoint, map_location="cpu") model.load_state_dict(ckpt["model_state_dict"]) model.to(device) + + # Restore event threshold from checkpoint, fall back to config default + event_threshold = ckpt.get("event_threshold", train_cfg.event_threshold) print(f"Loaded checkpoint from {args.checkpoint} (epoch={ckpt.get('epoch', '?')})") + print(f"Event threshold: {event_threshold} (checkpoint={ckpt.get('event_threshold', 'not saved')}, config={train_cfg.event_threshold})") # Evaluate each scene independently → NaN gaps prevent plot mixing from src.velocity_prediction.config import TEST_SCENES @@ -180,7 +184,7 @@ def main(): stride=1, batch_size=1, num_workers=0, # strict temporal order - event_threshold=train_cfg.event_threshold, + event_threshold=event_threshold, event_use_log=train_cfg.event_use_log, ) results = evaluate_stateful(model, loader, device) diff --git a/src/velocity_prediction/train.py b/src/velocity_prediction/train.py index 40f6ac0..0a39363 100644 --- a/src/velocity_prediction/train.py +++ b/src/velocity_prediction/train.py @@ -191,6 +191,10 @@ def main(): global_step = ckpt.get("global_step", 0) best_val_loss = ckpt.get("best_val_loss", float("inf")) run_id = ckpt.get("run_id", None) + # Restore event threshold from checkpoint if present + if "event_threshold" in ckpt: + event_threshold = ckpt["event_threshold"] + print(f"Event threshold restored from checkpoint: {event_threshold}") print(f"Resumed from checkpoint: {ckpt_path}") print(f" Resumed epoch={ckpt.get('epoch', '?')}, global_step={global_step}, " @@ -253,6 +257,7 @@ def main(): "model_state_dict": model.state_dict(), "optimizer_state_dict": optimizer.state_dict(), "scheduler_state_dict": scheduler.state_dict(), + "event_threshold": event_threshold, "global_step": global_step, "best_val_loss": best_val_loss, "run_id": run_id, @@ -269,6 +274,7 @@ def main(): "model_state_dict": model.state_dict(), "optimizer_state_dict": optimizer.state_dict(), "scheduler_state_dict": scheduler.state_dict(), + "event_threshold": event_threshold, "global_step": global_step, "best_val_loss": best_val_loss, "run_id": run_id,