| |
| 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() |
|
|
