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 <vibe@mistral.ai>
This commit is contained in:
2026-08-01 17:47:39 +08:00
parent ce70d932d3
commit a97b4da1ad
2 changed files with 11 additions and 1 deletions
+5 -1
View File
@@ -166,7 +166,11 @@ def main():
ckpt = torch.load(args.checkpoint, map_location="cpu") ckpt = torch.load(args.checkpoint, map_location="cpu")
model.load_state_dict(ckpt["model_state_dict"]) model.load_state_dict(ckpt["model_state_dict"])
model.to(device) 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"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 # Evaluate each scene independently → NaN gaps prevent plot mixing
from src.velocity_prediction.config import TEST_SCENES from src.velocity_prediction.config import TEST_SCENES
@@ -180,7 +184,7 @@ def main():
stride=1, stride=1,
batch_size=1, batch_size=1,
num_workers=0, # strict temporal order num_workers=0, # strict temporal order
event_threshold=train_cfg.event_threshold, event_threshold=event_threshold,
event_use_log=train_cfg.event_use_log, event_use_log=train_cfg.event_use_log,
) )
results = evaluate_stateful(model, loader, device) results = evaluate_stateful(model, loader, device)
+6
View File
@@ -191,6 +191,10 @@ def main():
global_step = ckpt.get("global_step", 0) global_step = ckpt.get("global_step", 0)
best_val_loss = ckpt.get("best_val_loss", float("inf")) best_val_loss = ckpt.get("best_val_loss", float("inf"))
run_id = ckpt.get("run_id", None) 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 from checkpoint: {ckpt_path}")
print(f" Resumed epoch={ckpt.get('epoch', '?')}, global_step={global_step}, " print(f" Resumed epoch={ckpt.get('epoch', '?')}, global_step={global_step}, "
@@ -253,6 +257,7 @@ def main():
"model_state_dict": model.state_dict(), "model_state_dict": model.state_dict(),
"optimizer_state_dict": optimizer.state_dict(), "optimizer_state_dict": optimizer.state_dict(),
"scheduler_state_dict": scheduler.state_dict(), "scheduler_state_dict": scheduler.state_dict(),
"event_threshold": event_threshold,
"global_step": global_step, "global_step": global_step,
"best_val_loss": best_val_loss, "best_val_loss": best_val_loss,
"run_id": run_id, "run_id": run_id,
@@ -269,6 +274,7 @@ def main():
"model_state_dict": model.state_dict(), "model_state_dict": model.state_dict(),
"optimizer_state_dict": optimizer.state_dict(), "optimizer_state_dict": optimizer.state_dict(),
"scheduler_state_dict": scheduler.state_dict(), "scheduler_state_dict": scheduler.state_dict(),
"event_threshold": event_threshold,
"global_step": global_step, "global_step": global_step,
"best_val_loss": best_val_loss, "best_val_loss": best_val_loss,
"run_id": run_id, "run_id": run_id,