File size: 2,408 Bytes
03573b6 | 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 | """Run complete-grid inference and preserve every sample in one NPZ."""
import sys
from pathlib import Path
import numpy as np
import torch
import yaml
ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
from model.precipitationsrcnn import FORMAT_VERSION, MODEL_NAME, bilinear_input, build_model
def load_checkpoint(path):
try:
return torch.load(path, map_location="cpu", weights_only=True)
except TypeError:
return torch.load(path, map_location="cpu")
def main():
config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
data = np.load(ROOT / config["data"]["root"] / "daily_precipitation.npz")
checkpoint = load_checkpoint(ROOT / config["paths"]["checkpoint"])
if checkpoint["format_version"] != FORMAT_VERSION or checkpoint["model_name"] != MODEL_NAME:
raise ValueError("checkpoint version/model mismatch")
if checkpoint["target_grid"] != [216, 488] or checkpoint["model_config"] != config["model"]:
raise ValueError("checkpoint shape/config mismatch")
model = build_model(checkpoint["model_config"])
model.load_state_dict(checkpoint["model"]); model.eval()
coarse = torch.from_numpy(data["coarse_precipitation"])
elevation = torch.from_numpy(np.repeat(data["elevation"], len(coarse), axis=0))
inputs = bilinear_input(coarse, elevation, (216, 488))
predictions = []
with torch.no_grad():
for index in range(len(inputs)):
predictions.append(model(inputs[index:index + 1]).numpy())
prediction = np.concatenate(predictions).astype(np.float32)
if prediction.shape != data["target_precipitation"].shape or not np.isfinite(prediction).all():
raise ValueError("incomplete or invalid inference output")
output = ROOT / config["paths"]["inference"]
output.parent.mkdir(parents=True, exist_ok=True)
np.savez_compressed(output, format_version=np.array(FORMAT_VERSION),
prediction=prediction, target=data["target_precipitation"],
bilinear_precipitation=inputs[:, :1].numpy(), elevation=data["elevation"],
coarse_precipitation=data["coarse_precipitation"],
timestamps=data["timestamps"], years=data["years"], units=data["units"])
print(f"predictions={output.relative_to(ROOT)} shape={prediction.shape}")
if __name__ == "__main__":
main()
|