File size: 4,962 Bytes
5c14105 | 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 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 | from __future__ import annotations
import argparse
import json
import sys
import time
from pathlib import Path
from typing import Any
import numpy as np
import torch
from torch.utils.data import DataLoader, Dataset
import yaml
PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
sys.path.insert(0, str(PROJECT_ROOT))
from model.neural_convection_fno import build_model_from_config
from scripts.train import load_pair, normalize_x, normalize_y, ssim_2d, tensor_to_numpy
class InferenceDataset(Dataset):
def __init__(self, x: torch.Tensor, y: torch.Tensor, stats: dict[str, torch.Tensor]) -> None:
self.x = x
self.y = y
self.stats = stats
def __len__(self) -> int:
return int(self.x.shape[0])
def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]:
return normalize_x(self.x[idx], self.stats), normalize_y(self.y[idx], self.stats)
def load_config(path: str | Path) -> dict[str, Any]:
with open(path, "r", encoding="utf-8") as handle:
return yaml.safe_load(handle)
def resolve(path: str | Path) -> Path:
path = Path(path)
return path if path.is_absolute() else PROJECT_ROOT / path
def select_device(name: str) -> torch.device:
if name == "auto":
return torch.device("cuda" if torch.cuda.is_available() else "cpu")
return torch.device(name)
def save_json(path: Path, payload: dict[str, Any]) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
with open(path, "w", encoding="utf-8") as handle:
json.dump(payload, handle, indent=2)
@torch.no_grad()
def main() -> None:
parser = argparse.ArgumentParser(description="Run compact FNO inference on a NeuralConvection split.")
parser.add_argument("--config", default=str(PROJECT_ROOT / "config/config.yaml"))
parser.add_argument("--split", default=None)
parser.add_argument("--checkpoint", default=None)
parser.add_argument("--device", default=None)
args = parser.parse_args()
config = load_config(args.config)
data_cfg = config["data"]
infer_cfg = config["inference"]
split = args.split or infer_cfg["split"]
device = select_device(args.device or config["experiment"].get("device", "auto"))
root = resolve(data_cfg["root_dir"])
stats = torch.load(resolve(data_cfg["stats_path"]), map_location="cpu")
x, y = load_pair(root, split, data_cfg["x_file"], data_cfg["y_file"])
loader = DataLoader(
InferenceDataset(x, y, stats),
batch_size=int(infer_cfg["batch_size"]),
shuffle=False,
num_workers=int(data_cfg.get("num_workers", 0)),
pin_memory=bool(data_cfg.get("pin_memory", True)) and device.type == "cuda",
)
model = build_model_from_config(config).to(device)
ckpt_path = resolve(args.checkpoint or infer_cfg["checkpoint_path"])
checkpoint = torch.load(ckpt_path, map_location=device)
model.load_state_dict(checkpoint["model_state_dict"])
model.eval()
preds: list[np.ndarray] = []
targets: list[np.ndarray] = []
inputs: list[np.ndarray] = []
start = time.perf_counter()
for x_batch, y_batch in loader:
x_batch = x_batch.to(device, non_blocking=True)
pred = tensor_to_numpy(model(x_batch))
preds.append(pred)
targets.append(tensor_to_numpy(y_batch))
inputs.append(tensor_to_numpy(x_batch))
elapsed = time.perf_counter() - start
pred_all = np.concatenate(preds, axis=0)
target_all = np.concatenate(targets, axis=0)
input_all = np.concatenate(inputs, axis=0)
channel_names = list(data_cfg["channel_names"])
metrics = {
"split": split,
"checkpoint_path": str(ckpt_path),
"num_samples": int(pred_all.shape[0]),
"time_seconds_per_sample": elapsed / max(1, pred_all.shape[0]),
"per_channel": {},
"paper_table5_compact_fno_2d": config["paper_table5_compact_fno_2d"],
}
for idx, name in enumerate(channel_names):
mse = float(np.mean((pred_all[:, idx] - target_all[:, idx]) ** 2))
ssim = float(np.mean([ssim_2d(pred_all[i, idx], target_all[i, idx]) for i in range(pred_all.shape[0])]))
metrics["per_channel"][name] = {"mse": mse, "ssim": ssim}
paper = config["paper_table5_compact_fno_2d"]["channels"].get(name, {})
print(
f"{name}: mse={mse:.8e} ssim={ssim:.6f} paper_mse={paper.get('mse')} paper_ssim={paper.get('ssim')}",
flush=True,
)
out_path = resolve(infer_cfg["output_path"])
out_path.parent.mkdir(parents=True, exist_ok=True)
np.savez_compressed(out_path, pred_norm=pred_all, target_norm=target_all, input_norm=input_all, channel_names=np.array(channel_names))
save_json(resolve(infer_cfg["metrics_path"]), metrics)
print(f"saved_predictions={out_path}", flush=True)
print(f"saved_metrics={resolve(infer_cfg['metrics_path'])}", flush=True)
if __name__ == "__main__":
main()
|