File size: 3,089 Bytes
eac1b4e 3cbde22 eac1b4e | 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 | 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()
|