File size: 2,226 Bytes
bd3493c | 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 | #!/usr/bin/env python
"""Run the one-day tas forecast through the official Aardvark modules."""
import argparse
import json
import sys
from pathlib import Path
ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
import copy
import torch
from model.aardvark_adapter import build_one_day_model, load_sample
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--root", type=Path, default=ROOT)
parser.add_argument("--device", default="cuda")
parser.add_argument("--checkpoint", type=Path, help="Optional checkpoint produced by scripts/train.py")
parser.add_argument("--output", type=Path, default=ROOT / "result" / "inference_one_day.json")
args = parser.parse_args()
root = args.root.resolve()
sample_path = root / "weights/sample_data/sample_data_final.pkl"
sample = load_sample(sample_path)
model = build_one_day_model(root / "weights", root / "official-src", args.device)
if args.checkpoint:
tuned = args.checkpoint.resolve()
payload = torch.load(tuned, map_location=args.device, weights_only=False)
target_model = model if payload["train_modules"] == "all" else model.sf_model
target_model.load_state_dict(payload["model"])
model.eval()
target = sample["y_target"].cpu()
with torch.inference_mode():
station, global_forecast, initial_state = model(copy.deepcopy(sample))
result_dir = root / "result"
result_dir.mkdir(parents=True, exist_ok=True)
torch.save(station.cpu(), result_dir / "prediction.pt")
torch.save(target, result_dir / "target.pt")
report = {
"device": args.device,
"lead_time_days": 1,
"station_tas_shape": list(station.shape),
"global_forecast_shape": list(global_forecast.shape),
"initial_state_shape": list(initial_state.shape),
"finite_outputs": bool(torch.isfinite(station).all()),
"weights": str(tuned) if args.checkpoint else "official",
}
args.output.parent.mkdir(parents=True, exist_ok=True)
args.output.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n")
print(json.dumps(report, indent=2, sort_keys=True))
if __name__ == "__main__":
main()
|