| from __future__ import annotations |
|
|
| import argparse |
| import sys |
| from pathlib import Path |
|
|
| SCRIPT_DIR = Path(__file__).resolve().parent |
| MODEL_DIR = SCRIPT_DIR.parent / "model" |
| for module_dir in (SCRIPT_DIR, MODEL_DIR): |
| if str(module_dir) not in sys.path: |
| sys.path.insert(0, str(module_dir)) |
|
|
| import numpy as np |
| import torch |
|
|
| from common import DEFAULT_CONFIG, active_model_config, load_config, resolve_path |
| from data import build_loader, load_statistics |
| from fourcastnet_v2 import ( |
| FourCastNetV2, |
| load_checkpoint, |
| ) |
|
|
|
|
| def choose_device() -> torch.device: |
| return torch.device("cuda", 0) if torch.cuda.is_available() else torch.device("cpu") |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser(description="Run FourCastNet v2 inference") |
| parser.add_argument("--config", default=str(DEFAULT_CONFIG)) |
| parser.add_argument("--checkpoint") |
| args = parser.parse_args() |
|
|
| config = load_config(args.config) |
| inference = config["inference"] |
|
|
| device = choose_device() |
| model = FourCastNetV2(active_model_config(config)).to(device) |
| if args.checkpoint: |
| checkpoint_path = resolve_path(config, args.checkpoint) |
| else: |
| checkpoint_path = resolve_path(config, inference["checkpoint_path"]) |
| result = load_checkpoint( |
| model, |
| checkpoint_path, |
| expected_profile=config["model"]["profile"], |
| expected_variables=config["data"]["variables"], |
| allowed_stages={"one_step", "finetune"}, |
| allowed_initializations={"random", "one_step_checkpoint"}, |
| strict=config["checkpoint"]["strict"], |
| map_location=device, |
| ) |
| if result["missing_keys"] or result["unexpected_keys"]: |
| print( |
| f"missing_keys={result['missing_keys']} " |
| f"unexpected_keys={result['unexpected_keys']}" |
| ) |
|
|
| loader, _ = build_loader( |
| config, |
| config["data"]["test_years"], |
| train=False, |
| distributed=False, |
| output_steps=inference["rollout_steps"], |
| ) |
| means, stds = load_statistics(config) |
| means = means.numpy() |
| stds = stds.numpy() |
|
|
| output_dir = resolve_path(config, inference["output_dir"]) |
| output_dir.mkdir(parents=True, exist_ok=True) |
| model.eval() |
| with torch.no_grad(): |
| for sample_index, batch in enumerate(loader): |
| if sample_index >= inference["max_samples"]: |
| break |
| inputs = batch[0].to(device) |
| targets = batch[1] |
| state = inputs |
| predictions = [] |
| for _ in range(inference["rollout_steps"]): |
| state = model(state) |
| predictions.append(state.cpu()) |
| prediction = torch.stack(predictions, dim=1).numpy() |
| if targets.ndim == 4: |
| targets = targets.unsqueeze(1) |
| target = targets.numpy() |
| input_array = inputs.cpu().numpy() |
| if not inference["save_normalized"]: |
| prediction = prediction * stds[:, None] + means[:, None] |
| target = target * stds[:, None] + means[:, None] |
| input_array = input_array * stds + means |
| path = output_dir / f"sample_{sample_index:04d}.npz" |
| np.savez_compressed( |
| path, |
| input=input_array, |
| prediction=prediction, |
| target=target, |
| variables=np.asarray(config["data"]["variables"]), |
| time_index=np.asarray(batch[4], dtype=str).T, |
| ) |
| print(path) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
