| from __future__ import annotations |
|
|
| import argparse |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
|
|
| from common import ( |
| ROOT, |
| UnitTransformer, |
| build_model, |
| load_config, |
| load_pipe_tensors, |
| relative_l2_per_sample, |
| resolve, |
| save_json, |
| select_device, |
| ) |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser(description="Run HO-FNO Pipe inference with the best checkpoint.") |
| parser.add_argument("--config", default=str(ROOT / "config" / "config.yaml")) |
| parser.add_argument("--checkpoint", default=None) |
| args = parser.parse_args() |
|
|
| cfg = load_config(args.config) |
| output_dir = resolve(cfg["paths"]["output_dir"]) |
| results_dir = resolve(cfg["paths"]["results_dir"], output_dir) |
| results_dir.mkdir(parents=True, exist_ok=True) |
| checkpoint = resolve(args.checkpoint or cfg["paths"]["checkpoint"], output_dir) |
| device = select_device(require_accelerator=False) |
|
|
| payload = torch.load(checkpoint, map_location=device) |
| ckpt_cfg = payload.get("config", cfg) |
| ckpt_cfg["paths"]["data_dir"] = cfg["paths"]["data_dir"] |
| model = build_model(ckpt_cfg).to(device) |
| model.load_state_dict(payload["model_state"]) |
| model.eval() |
| x_norm = UnitTransformer.from_state_dict(payload["x_normalizer"]).to(device) |
| y_norm = UnitTransformer.from_state_dict(payload["y_normalizer"]).to(device) |
|
|
| _, _, x_test, y_test = load_pipe_tensors(ckpt_cfg) |
| batch_size = int(ckpt_cfg["training"]["batch_size"]) |
| predictions = [] |
| targets = [] |
| inputs = [] |
| rel_values = [] |
| with torch.no_grad(): |
| for start in range(0, x_test.shape[0], batch_size): |
| x = x_test[start : start + batch_size].to(device) |
| y = y_test[start : start + batch_size].to(device) |
| pred = y_norm.decode(model(x_norm.encode(x))) |
| rel = relative_l2_per_sample(pred, y) |
| predictions.append(pred.cpu().numpy()) |
| targets.append(y.cpu().numpy()) |
| inputs.append(x.cpu().numpy()) |
| rel_values.append(rel.cpu().numpy()) |
|
|
| pred_arr = np.concatenate(predictions, axis=0) |
| target_arr = np.concatenate(targets, axis=0) |
| input_arr = np.concatenate(inputs, axis=0) |
| rel_arr = np.concatenate(rel_values, axis=0) |
| np.save(results_dir / "predictions.npy", pred_arr) |
| np.save(results_dir / "targets.npy", target_arr) |
| np.save(results_dir / "inputs.npy", input_arr) |
| np.save(results_dir / "relative_l2_per_sample.npy", rel_arr) |
| metrics = { |
| "checkpoint": str(checkpoint), |
| "test_samples": int(rel_arr.shape[0]), |
| "mean_relative_l2": float(rel_arr.mean()), |
| "median_relative_l2": float(np.median(rel_arr)), |
| "max_relative_l2": float(rel_arr.max()), |
| "min_relative_l2": float(rel_arr.min()), |
| } |
| save_json(results_dir / "inference_metrics.json", metrics) |
| print( |
| "inference_complete " |
| f"mean_relative_l2={metrics['mean_relative_l2']:.8f} " |
| f"predictions={results_dir / 'predictions.npy'}", |
| flush=True, |
| ) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
