File size: 4,283 Bytes
9f29df6 | 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 100 101 102 103 104 105 106 107 108 | """Checkpoint-backed one-step and autoregressive RainNet inference."""
import json
from pathlib import Path
import sys
import h5py
import numpy as np
import torch
import torch.nn.functional as F
import yaml
ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
from model.rainnet import build_rainnet
def load_config():
with (ROOT / "conf/config.yaml").open(encoding="utf-8") as handle:
return yaml.safe_load(handle)
def inverse_transform(tensor):
return torch.clamp(torch.exp(tensor) - 0.01, min=0.0)
def describe(name, array):
print(
f"{name}: shape={array.shape}, dtype={array.dtype}, "
f"min={array.min():.8f}, max={array.max():.8f}, mean={array.mean():.8f}"
)
def main():
config = load_config()
torch.manual_seed(config["seed"])
requested_device = config["device"]
if requested_device == "auto":
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
else:
device = torch.device(requested_device)
checkpoint_path = ROOT / config["inference"]["checkpoint"]
if not checkpoint_path.exists():
raise FileNotFoundError(f"Required checkpoint does not exist: {checkpoint_path}")
model = build_rainnet(**config["model"]).to(device)
checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
model.load_state_dict(checkpoint["model_state_dict"])
model.eval()
print("checkpoint_loaded: True")
print(f"checkpoint path: {checkpoint_path}")
data = config["data"]
data_path = ROOT / data["path"]
with h5py.File(data_path, "r") as handle:
keys = sorted(handle.keys())
test_start = data["train_frames"] + data["val_frames"]
test_keys = keys[test_start : test_start + data["test_frames"]]
raw = np.stack([handle[key][...] for key in test_keys]).astype(np.float32)
needed = data["input_steps"] + config["inference"]["rollout_steps"]
if len(raw) < needed:
raise ValueError(f"Inference needs {needed} test frames, found {len(raw)}")
inputs = torch.from_numpy(raw[: data["input_steps"]]).unsqueeze(0).to(device)
targets = raw[data["input_steps"] : needed]
log_window = torch.log(inputs + 0.01)
pad_h = data["padded_height"] - data["raw_height"]
pad_w = data["padded_width"] - data["raw_width"]
pad = (pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2)
crop = (pad_h // 2, pad_w // 2)
predictions = []
with torch.inference_mode():
for _ in range(config["inference"]["rollout_steps"]):
padded = F.pad(log_window, pad, mode="reflect")
padded_prediction = model(padded)
prediction = padded_prediction[
:, :, crop[0] : crop[0] + data["raw_height"], crop[1] : crop[1] + data["raw_width"]
]
predictions.append(inverse_transform(prediction).squeeze(0).squeeze(0).cpu().numpy())
log_window = torch.cat((log_window[:, 1:], prediction), dim=1)
predictions = np.stack(predictions).astype(np.float32)
persistence = np.repeat(raw[data["input_steps"] - 1][None], len(predictions), axis=0).astype(np.float32)
output_dir = ROOT / config["inference"]["output_dir"]
output_dir.mkdir(parents=True, exist_ok=True)
np.save(output_dir / "inputs.npy", raw[: data["input_steps"]])
np.save(output_dir / "predictions.npy", predictions)
np.save(output_dir / "targets.npy", targets)
np.save(output_dir / "persistence.npy", persistence)
metadata = {
"units": "mm/5min",
"interval_minutes": data["interval_minutes"],
"rollout_steps": len(predictions),
"input_keys": test_keys[: data["input_steps"]],
"target_keys": test_keys[data["input_steps"] : needed],
"checkpoint": str(checkpoint_path),
}
with (output_dir / "metadata.json").open("w", encoding="utf-8") as handle:
json.dump(metadata, handle, indent=2)
describe("inputs", raw[: data["input_steps"]])
describe("predictions", predictions)
describe("targets", targets)
describe("persistence", persistence)
print(f"One-step inference shape: {predictions[:1].shape}")
print(f"Autoregressive rollout steps: {len(predictions)}")
if __name__ == "__main__":
main()
|