File size: 3,330 Bytes
30f7852 | 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 | """Apply all four UnetDif heads and save 3-hour and 24-hour products."""
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.unetdif import UnetDif, apply_correction
def main():
config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
checkpoint = torch.load(ROOT / config["paths"]["checkpoint"], map_location=device, weights_only=True)
required_checkpoint = {"epoch", "model", "optimizer", "model_config", "loss_normalizers", "loss_components", "format_version"}
if not required_checkpoint.issubset(checkpoint):
raise ValueError(f"checkpoint is missing {sorted(required_checkpoint - checkpoint.keys())}")
if checkpoint["format_version"] != config["data"]["format_version"]:
raise ValueError("checkpoint and configured data format versions differ")
model = UnetDif(**checkpoint["model_config"]).to(device)
model.load_state_dict(checkpoint["model"]); model.eval()
archive = np.load(ROOT / config["data"]["root"] / "test.npz", allow_pickle=False)
inputs = archive["inputs"].astype(np.float32, copy=False)
expected = (int(config["data"]["time_steps"]), int(config["data"]["channels"]),
int(config["data"]["height"]), int(config["data"]["width"]))
if inputs.shape[1:] != expected or archive["targets"].shape[1:] != (expected[0], expected[2], expected[3]):
raise ValueError(f"test NPZ does not satisfy the fixed day contract {expected}")
if archive["lead_hours"].tolist() != config["data"]["lead_hours"]:
raise ValueError("test NPZ lead_hours do not match configuration")
flat = torch.from_numpy(inputs.reshape(-1, *inputs.shape[2:])).to(device)
heads_all = {key: [] for key in ("dry_logits", "false_alarm_logits", "positive_bias", "negative_bias")}
corrected = []
batch_size = int(config["inference"]["batch_size"])
with torch.no_grad():
for start in range(0, len(flat), batch_size):
batch = flat[start:start + batch_size]
heads = model(batch)
corrected.append(apply_correction(batch[:, int(config["data"]["precipitation_channel"])], heads,
float(config["inference"]["probability_threshold"])).cpu())
for key in heads_all:
heads_all[key].append(heads[key].cpu())
day_shape = inputs.shape[:2] + inputs.shape[-2:]
corrected_3h = torch.cat(corrected).numpy().reshape(day_shape)
raw_3h = inputs[:, :, int(config["data"]["precipitation_channel"])]
output = ROOT / config["paths"]["inference"]
output.parent.mkdir(parents=True, exist_ok=True)
payload = {key: torch.cat(value).numpy().reshape(day_shape) for key, value in heads_all.items()}
np.savez_compressed(
output, inputs=inputs, group_id=archive["group_id"], lead_hours=archive["lead_hours"],
raw_3h=raw_3h, target_3h=archive["targets"], corrected_3h=corrected_3h,
raw_24h=raw_3h.sum(1), target_24h=archive["targets"].sum(1), corrected_24h=corrected_3h.sum(1), **payload,
)
print(f"saved={output.relative_to(ROOT)} corrected_3h={corrected_3h.shape}")
if __name__ == "__main__":
main()
|