File size: 1,234 Bytes
38f9214 | 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 | """Load the compact MetNet-3 checkpoint and run fake-data inference."""
import json
import sys
from pathlib import Path
import torch
ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
sys.path.insert(0, str(Path(__file__).resolve().parent))
from model import MetNet3, MetNet3Config
from model.fake_data import make_fake
def main() -> None:
result_dir = ROOT / "result"
result_dir.mkdir(parents=True, exist_ok=True)
checkpoint = torch.load(ROOT / "weight" / "model.pth", map_location="cpu", weights_only=True)
config = MetNet3Config(**checkpoint["config"])
model = MetNet3(config)
model.load_state_dict(checkpoint["model"])
model.eval()
batch, targets = make_fake(config)
with torch.inference_mode():
outputs = model(batch)
torch.save(outputs, result_dir / "prediction.pt")
torch.save(targets, result_dir / "target.pt")
summary = {name: list(value.shape) for name, value in outputs.items()}
summary["finite"] = all(bool(torch.isfinite(value).all()) for value in outputs.values())
(result_dir / "inference.json").write_text(json.dumps(summary, indent=2) + "\n")
print(json.dumps(summary, indent=2))
if __name__ == "__main__":
main()
|