File size: 2,627 Bytes
e5a6e71 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 | """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()
|