| """Reload a trained checkpoint; stream evaluation from the supplied MATLAB HDF5.""" |
| import argparse |
| import json |
| import sys |
| from pathlib import Path |
| sys.path.insert(0, str(Path(__file__).resolve().parents[1])) |
| import h5py |
| import numpy as np |
| import torch |
| from common import (ROOT, atomic_json, build_model, load_checkpoint, load_config, |
| predict, relative_metrics, resolved, seed_everything, sha256, |
| validation_indices) |
|
|
|
|
| def infer(config, checkpoint_path, output_dir, metrics_only=False): |
| seed_everything(0, config["training"]["num_threads"]) |
| model = build_model(config) |
| checkpoint = load_checkpoint(checkpoint_path, model) |
| if checkpoint["config"]["model"] != config["model"]: |
| raise ValueError("Checkpoint and requested model configuration differ") |
| model.eval() |
| out = Path(output_dir) |
| out.mkdir(parents=True, exist_ok=True) |
| source_path = resolved(config, "reference") |
| report = {"checkpoint": str(Path(checkpoint_path).resolve()), "checkpoint_sha256": sha256(checkpoint_path), |
| "source": str(source_path), "source_sha256": sha256(source_path), |
| "seed": checkpoint["seed"], "epoch": checkpoint["epoch"], "selection_score": checkpoint["selection_score"], |
| "validation_metrics": checkpoint["validation_metrics"], |
| "metric": "raw ||pred-ref||2/||ref||2; no pressure gauge alignment", |
| "groups": {}, "reference_regenerated": False, "checkpoint_reloaded": True} |
| size = config["inference"]["batch_size"] |
| global_num, global_den = np.zeros(3), np.zeros(3) |
| output = None if metrics_only else h5py.File(out/"predictions.h5.tmp", "w") |
| try: |
| with h5py.File(source_path, "r") as source: |
| for key in config["inference"]["reference_keys"]: |
| data = source[key] |
| if data.ndim != 2 or data.shape[0] != 6: raise ValueError(f"Invalid {key} shape: {data.shape}") |
| count = data.shape[1] |
| numer, denom, hold_num, hold_den = [np.zeros(3) for _ in range(4)] |
| if key == config["validation"]["reference_key"]: |
| valids = validation_indices(config, count) |
| group = None |
| if output is not None: |
| group = output.create_group(key) |
| group.create_dataset("coordinates", (count, 3), dtype="f8", chunks=True) |
| group.create_dataset("prediction", (count, 3), dtype="f4", chunks=True) |
| for start in range(0, count, size): |
| end = min(start+size, count) |
| rows = data[:, start:end].T |
| if not np.isfinite(rows).all(): raise ValueError(f"Nonfinite input in {key} at {start}") |
| preds = predict(model, rows[:, :3], size) |
| if not np.isfinite(preds).all(): raise FloatingPointError("Nonfinite predictions") |
| error2 = (preds.astype(np.float64)-rows[:, 3:])**2 |
| ref2 = rows[:, 3:]**2 |
| numer += error2.sum(0) |
| denom += ref2.sum(0) |
| if key == config["validation"]["reference_key"]: |
| keep = ~np.isin(np.arange(start, end), valids, assume_unique=True) |
| hold_num += error2[keep].sum(0) |
| hold_den += ref2[keep].sum(0) |
| if group is not None: |
| group["coordinates"][start:end] = rows[:, :3] |
| group["prediction"][start:end] = preds |
| metrics = relative_metrics(numer, denom) |
| report["groups"][key] = {"count": count, "metrics": metrics} |
| global_num += numer |
| global_den += denom |
| if key == config["validation"]["reference_key"]: |
| report["internal_holdout"] = {"count": count-len(valids), "metrics": relative_metrics(hold_num, hold_den)} |
| print(f"INFERENCE key={key} count={count} metrics={json.dumps(metrics)}", flush=True) |
| report["all_reference"] = {"metrics": relative_metrics(global_num, global_den), |
| "count": sum(v["count"] for v in report["groups"].values())} |
| if output is not None: |
| output.attrs["provenance"] = json.dumps({k: v for k, v in report.items() if k not in ["groups", "internal_holdout", "all_reference"]}) |
| output.close() |
| output = None |
| (out/"predictions.h5.tmp").replace(out/"predictions.h5") |
| atomic_json(out/"metrics.json", report) |
| finally: |
| if output is not None: output.close() |
| print("INFERENCE_COMPLETED", flush=True) |
| return report |
|
|
|
|
| if __name__ == "__main__": |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--config", default=str(ROOT/"config/config.yaml")) |
| parser.add_argument("--checkpoint", type=Path) |
| parser.add_argument("--output-dir", type=Path) |
| parser.add_argument("--seed", type=int) |
| parser.add_argument("--kind", choices=["best", "last"], default="best") |
| parser.add_argument("--metrics-only", action="store_true") |
| args = parser.parse_args() |
| cfg = load_config(args.config) |
| checkpoint = args.checkpoint or (resolved(cfg, "weights")/f"seed_{args.seed}"/("best_model.pt" if args.kind == "best" else "last_model.pt") if args.seed is not None else resolved(cfg, "weights")/"best_model.pt") |
| output = args.output_dir or (resolved(cfg, "results")/f"seed_{args.seed}"/("best_eval" if args.kind == "best" else "final_eval") if args.seed is not None else resolved(cfg, "results")/"best_model") |
| infer(cfg, checkpoint, output, args.metrics_only) |
|
|
