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