File size: 3,283 Bytes
bdccc62 | 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 | from __future__ import annotations
import argparse
import sys
from pathlib import Path
import numpy as np
import torch
PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
sys.path.insert(0, str(PROJECT_ROOT))
from model.ContinuousFlowOperator import ContinuousFlowOperator, rmse, trajectory_relative_l2 # noqa: E402
from scripts.common import ensure_dirs, load_burgers_splits, load_config, write_json # noqa: E402
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Run full-test inference for CFO Burgers")
parser.add_argument("--config", default=str(PROJECT_ROOT / "config" / "config.yaml"))
parser.add_argument("--checkpoint", default=None)
return parser.parse_args()
@torch.no_grad()
def rollout_all(model: ContinuousFlowOperator, data: np.ndarray, device: torch.device, steps_per_segment: int, batch_size: int) -> np.ndarray:
preds = []
for start in range(0, len(data), batch_size):
x0 = torch.from_numpy(data[start : start + batch_size, 0]).to(device=device, dtype=torch.float32)
preds.append(model.rollout(x0, data.shape[1], steps_per_segment=steps_per_segment).numpy())
return np.concatenate(preds, axis=0)
def main() -> None:
args = parse_args()
cfg = load_config(args.config)
ensure_dirs(cfg)
ckpt_path = Path(args.checkpoint or cfg["inference"]["checkpoint"])
if not ckpt_path.exists():
raise FileNotFoundError(f"checkpoint not found: {ckpt_path}")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
splits, tspan = load_burgers_splits(cfg)
test_data = splits["test"]
model = ContinuousFlowOperator(input_size=cfg["model"]["input_size"], gamma=cfg["model"]["gamma"]).to(device)
ckpt = torch.load(ckpt_path, map_location=device)
model.load_state_dict(ckpt["model_state_dict"])
model.eval()
pred = rollout_all(
model,
test_data,
device,
steps_per_segment=int(cfg["inference"]["steps_per_segment"]),
batch_size=int(cfg["inference"]["eval_batch_size"]),
)
per_sample = []
for i in range(len(test_data)):
per_sample.append(trajectory_relative_l2(test_data[i : i + 1], pred[i : i + 1]))
metrics = {
"checkpoint": str(ckpt_path),
"test_shape": list(test_data.shape),
"prediction_shape": list(pred.shape),
"trajectory_relative_l2": trajectory_relative_l2(test_data, pred),
"rmse": rmse(test_data, pred),
"per_sample_relative_l2_mean": float(np.mean(per_sample)),
"per_sample_relative_l2_std": float(np.std(per_sample)),
"per_sample_relative_l2_max": float(np.max(per_sample)),
"worst_sample_index": int(np.argmax(per_sample)),
"solver": cfg["inference"]["solver"],
"steps_per_segment": int(cfg["inference"]["steps_per_segment"]),
}
out_dir = Path(cfg["paths"]["results_dir"])
np.savez_compressed(out_dir / "test_predictions.npz", prediction=pred, target=test_data, tspan=tspan, per_sample_relative_l2=np.asarray(per_sample))
write_json(out_dir / "test_metrics.json", metrics)
print("Full test inference complete.", flush=True)
print(metrics, flush=True)
if __name__ == "__main__":
main()
|