| """Reload the selected checkpoint and evaluate every short/long test trajectory.""" |
| import json |
| from pathlib import Path |
|
|
| import torch |
|
|
| from common import (ROOT, Trajectories, build_model, checkpoint_config_matches, device, load_config, output_dir, |
| parser, provenance, rollout_metrics, save_json, setup, sha256) |
|
|
|
|
| def main(): |
| arg = parser(__doc__) |
| arg.add_argument("--checkpoint", type=Path) |
| args = arg.parse_args() |
| cfg = load_config(args.config) |
| setup(cfg["seed"]) |
| target_device = device() |
| out = output_dir(cfg) |
| checkpoint_path = args.checkpoint or ROOT / cfg["paths"]["best_checkpoint"] |
| checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False) |
| status = json.loads((out / "training_status.json").read_text()) |
| if status["state"] != "completed" or status["completed_epochs"] != 300: |
| raise RuntimeError("Final inference requires evidence of completed 300-epoch training") |
| if not checkpoint_config_matches(checkpoint["config"], cfg) or checkpoint["epoch"] != status["best_epoch"]: |
| raise ValueError("Checkpoint/config/best-epoch provenance mismatch") |
| model = build_model(cfg, target_device) |
| model.load_state_dict(checkpoint["model_state_dict"], strict=True) |
| if not all(torch.isfinite(p).all() for p in model.parameters()): |
| raise ValueError("Nonfinite checkpoint parameters") |
| information = {"checkpoint": str(checkpoint_path), "sha256": sha256(checkpoint_path), |
| "selected_epoch": checkpoint["epoch"], "completed_epochs": 300, |
| "best_val_mse": checkpoint["best_val_loss"], |
| "parameter_count": sum(p.numel() for p in model.parameters()), |
| "provenance": provenance()} |
| print("CHECKPOINT_VERIFIED " + json.dumps(information), flush=True) |
| metrics, manifests = {}, {} |
| for split in cfg["evaluation"]["splits"]: |
| data = Trajectories(cfg, split) |
| manifests[split] = data.metadata |
| metrics[split] = rollout_metrics(model, data, target_device, |
| batch_size=cfg["evaluation"]["batch_size"], |
| save_path=out / split, |
| example_cases=cfg["evaluation"]["example_cases"]) |
| print("FINAL_METRICS " + json.dumps(metrics[split]), flush=True) |
| save_json(out / "inference_manifest.json", {**information, "data": manifests}) |
| save_json(out / "metrics.json", metrics) |
| print("INFERENCE_COMPLETED all=100 trajectories long_steps=120", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
