File size: 1,844 Bytes
8791cfe | 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 | #!/usr/bin/env python3
import argparse
from pathlib import Path
import numpy as np
import torch
from model.nncam import NNCAM, unscale_output
ROOT = Path(__file__).resolve().parents[1]
def main():
parser = argparse.ArgumentParser(description="Run offline NNCAM inference.")
parser.add_argument("--data", type=Path, default=ROOT / "data/nncam_fake.npz")
parser.add_argument("--checkpoint", type=Path, default=ROOT / "result/checkpoints/nncam.pt")
parser.add_argument("--output", type=Path, default=ROOT / "result/output/predictions.npz")
args = parser.parse_args()
checkpoint = torch.load(args.checkpoint, map_location="cpu", weights_only=True)
required = {"model", "model_config", "format_version", "normalization"}
if not required.issubset(checkpoint):
raise ValueError(f"checkpoint missing {sorted(required - checkpoint.keys())}")
model = NNCAM(**checkpoint["model_config"])
model.load_state_dict(checkpoint["model"])
model.eval()
with np.load(args.data) as data:
x, truth, lat, time = (data[name] for name in ("x", "y", "lat", "time"))
norm = checkpoint["normalization"]
normalized = (torch.from_numpy(x) - norm["input_mean"]) / norm["input_scale"]
with torch.no_grad():
scaled = model(normalized) * norm["target_scale"] + norm["target_mean"]
prediction = unscale_output(scaled.numpy()).astype(np.float32)
if prediction.shape != truth.shape or not np.isfinite(prediction).all():
raise RuntimeError(f"invalid prediction shape or values: {prediction.shape}")
args.output.parent.mkdir(parents=True, exist_ok=True)
np.savez_compressed(args.output, input=x, truth=truth, prediction=prediction, lat=lat, time=time)
print(f"saved {args.output}: prediction={prediction.shape}, finite=true")
if __name__ == "__main__":
main()
|