File size: 5,712 Bytes
d647964 | 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 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 | """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)
|