File size: 2,047 Bytes
73d9e73 | 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 | """Load an SFNO smoke checkpoint and run autoregressive inference."""
import argparse
import json
import sys
from pathlib import Path
import torch
ROOT = Path(__file__).resolve().parents[1]
LOCAL_DEPS = ROOT / ".deps"
if LOCAL_DEPS.is_dir():
sys.path.insert(0, str(LOCAL_DEPS))
sys.path.insert(0, str(ROOT))
from model.config import load_config
from model.fake_spherical_data import make_fake_spherical_sequence
from model.sfno_adapter import OfficialSFNOAdapter
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--checkpoint", type=Path, default=ROOT / "weight" / "model.pth")
parser.add_argument("--output-dir", type=Path, default=ROOT / "result")
args = parser.parse_args()
result_dir = args.output_dir
result_dir.mkdir(parents=True, exist_ok=True)
checkpoint = torch.load(args.checkpoint, map_location="cpu", weights_only=True)
config_path = result_dir / "checkpoint_config.json"
config_path.write_text(json.dumps(checkpoint["config"], indent=2) + "\n")
config = load_config(config_path)
model = OfficialSFNOAdapter(config)
model.load_state_dict(checkpoint["model"])
model.eval()
fields = make_fake_spherical_sequence(
config.timesteps, config.channels, config.nlat, config.nlon, config.seed
)["fields"]
state = fields[0:1]
target = fields[1 : config.rollout_steps + 1]
forecasts = []
with torch.inference_mode():
for _ in range(config.rollout_steps):
state = model(state)
forecasts.append(state)
prediction = torch.cat(forecasts, dim=0)
torch.save(prediction, result_dir / "prediction.pt")
torch.save(target, result_dir / "target.pt")
summary = {
"prediction_shape": list(prediction.shape),
"target_shape": list(target.shape),
"finite": bool(torch.isfinite(prediction).all()),
}
(result_dir / "inference.json").write_text(json.dumps(summary, indent=2) + "\n")
print(json.dumps(summary, indent=2))
if __name__ == "__main__":
main()
|