File size: 4,560 Bytes
f637d77 | 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 |
import argparse
from pathlib import Path
import importlib.util
import numpy as np
import onnxruntime as ort
# -----------------------------
# Helpers: import modules by file path (avoid "utils" name conflicts)
# -----------------------------
ROOT = Path(__file__).resolve().parents[1] # inference_example/
def _import_from_path(module_name: str, file_path: Path):
spec = importlib.util.spec_from_file_location(module_name, str(file_path))
if spec is None or spec.loader is None:
raise ImportError(f"Cannot import {module_name} from {file_path}")
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
return module
# Load our local helper files explicitly
io = _import_from_path("fuxicfd_io", ROOT / "utils" / "io.py")
pre = _import_from_path("fuxicfd_pre", ROOT / "utils" / "preprocessing.py")
post = _import_from_path("fuxicfd_post", ROOT / "utils" / "postprocessing.py")
load_example_input = io.load_example_input
save_prediction_npz = io.save_prediction_npz
build_model_input = pre.build_model_input
denormalize_and_split = post.denormalize_and_split
def _resolve_path(p: str) -> Path:
"""Resolve a path relative to inference_example/ if it's not absolute."""
path = Path(p)
if path.is_absolute():
return path
return (ROOT / path).resolve()
def main():
parser = argparse.ArgumentParser(description="FuXi-CFD ONNX inference example (publish-ready).")
parser.add_argument("--model", type=str, default="../model/fuxicfd_model.onnx",
help="Path to ONNX model. Default: ../model/fuxicfd_model.onnx (relative to inference_example/)")
parser.add_argument("--input", type=str, default="data/inputs.npz",
help="Input file (.npy dict or .npz). Default: data/inputs.npz")
parser.add_argument("--output", type=str, default="data/prediction.npz",
help="Output prediction .npz. Default: data/prediction.npz")
parser.add_argument("--norm_in", type=str, default="normalization/scaler_input.npy",
help="Input normalization stats (.npy dict). Default: normalization/scaler_input.npy")
parser.add_argument("--norm_out", type=str, default="normalization/scaler_output.npy",
help="Output normalization stats (.npy dict). Default: normalization/scaler_output.npy")
parser.add_argument("--device", type=str, default="cpu", choices=["cpu", "cuda"],
help="Execution device for onnxruntime. Default: cpu")
args = parser.parse_args()
# Resolve paths relative to inference_example/
model_path = _resolve_path(args.model)
input_path = _resolve_path(args.input)
output_path = _resolve_path(args.output)
norm_in_path = _resolve_path(args.norm_in)
norm_out_path = _resolve_path(args.norm_out)
# --- sanity checks ---
if not model_path.exists():
raise FileNotFoundError(
f"ONNX model file not found:\n {model_path}\n\n"
f"Please place your model at:\n {ROOT.parent / 'model' / 'fuxicfd_model.onnx'}\n"
f"or pass --model with the correct path."
)
if not norm_in_path.exists():
raise FileNotFoundError(f"Input normalization file not found: {norm_in_path}")
if not norm_out_path.exists():
raise FileNotFoundError(f"Output normalization file not found: {norm_out_path}")
if not input_path.exists():
raise FileNotFoundError(f"Input file not found: {input_path}")
# Load normalization dicts
in_stats = np.load(norm_in_path, allow_pickle=True).item()
out_stats = np.load(norm_out_path, allow_pickle=True).item()
# Select providers
providers = ["CPUExecutionProvider"]
if args.device == "cuda":
providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
# Create session
sess = ort.InferenceSession(str(model_path), providers=providers)
input_name = sess.get_inputs()[0].name
# Load and preprocess inputs
x_dict = load_example_input(str(input_path))
x = build_model_input(x_dict, in_stats) # (4,300,300) float32
ort_inputs = x[None, ...] # (1,4,300,300)
# Run inference
pred = sess.run(None, {input_name: ort_inputs})[0] # (1,27,4,300,300) expected
# Postprocess
u, v, w, k = denormalize_and_split(pred, out_stats)
# Save output
save_prediction_npz(str(output_path), u=u, v=v, w=w, k=k)
print(f"[OK] Saved prediction to: {output_path}")
if __name__ == "__main__":
main()
|