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:
@@ -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)
|
||||||
|
|||||||
@@ -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,
|
||||||
|
|||||||
Reference in New Issue
Block a user