File size: 2,523 Bytes
939693f | 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 | #!/usr/bin/env python3
"""Smoke-test FireWx-FM checkpoint loading and tensor inference.
This test intentionally uses a synthetic 16-channel tensor. It verifies that
the released model code and checkpoint are mutually compatible; scientific
reproduction still requires the source datasets documented in data_sources/.
"""
from __future__ import annotations
import argparse
import sys
from pathlib import Path
import torch
ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
sys.path.insert(0, str(ROOT))
from firewxfm.modeling_unet import UNetSmallFlex
from firewxfm.serve_conus import load_checkpoint
def main() -> int:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument(
"--checkpoint",
type=Path,
default=Path("models/checkpoints/firewxfm_2024_staticfix_region_balanced_bce_seed42.pt"),
)
parser.add_argument("--height", type=int, default=64)
parser.add_argument("--width", type=int, default=64)
parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
args = parser.parse_args()
device = torch.device(args.device)
torch.manual_seed(0)
model = UNetSmallFlex(
in_ch=16,
base=32,
dropout=0.1,
norm_type="group",
norm_groups=8,
use_aux_spatial_head=True,
).to(device)
state = load_checkpoint(args.checkpoint, device)
missing, unexpected = model.load_state_dict(state, strict=False)
model.eval()
x = torch.randn(1, 16, args.height, args.width, device=device)
with torch.no_grad():
logits = model(x)
prob = torch.sigmoid(logits)
expected = (1, 1, args.height, args.width)
if tuple(prob.shape) != expected:
raise RuntimeError(f"Expected output shape {expected}, got {tuple(prob.shape)}")
if not torch.isfinite(prob).all():
raise RuntimeError("Output contains non-finite values.")
if float(prob.min()) < 0.0 or float(prob.max()) > 1.0:
raise RuntimeError("Sigmoid output is outside [0, 1].")
print("FireWx-FM smoke inference passed.")
print(f"checkpoint={args.checkpoint}")
print(f"device={device}")
print(f"output_shape={tuple(prob.shape)}")
print(f"prob_mean={float(prob.mean()):.6f} prob_max={float(prob.max()):.6f}")
if missing:
print(f"missing_keys={len(missing)}")
if unexpected:
print(f"unexpected_keys={len(unexpected)}")
return 0
if __name__ == "__main__":
raise SystemExit(main())
|