File size: 2,997 Bytes
6c3f19f | 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 | import argparse
import json
import os
import sys
from pathlib import Path
import numpy as np
import torch
PROJECT_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(PROJECT_ROOT))
from scripts.common import (
PROBLEMS,
build_models,
ensure_onescience_path,
load_checkpoint,
load_config,
resolve_path,
select_device,
tensor_to_numpy_dict,
)
ensure_onescience_path()
from scripts.topology_optimization import clear_cached_kernels, predict_fields, share_mean_module
def parse_args():
parser = argparse.ArgumentParser(description="Run GP_for_TO field inference from a checkpoint.")
parser.add_argument("--problem", choices=PROBLEMS, default=None)
parser.add_argument("--gpu", type=int, default=None)
parser.add_argument("--device", default=None)
parser.add_argument("--n-col-domain", type=int, default=None)
parser.add_argument("--n-train-per-bc", type=int, default=None)
parser.add_argument("--checkpoint-path", default=None)
parser.add_argument("--output-dir", default=None)
return parser.parse_args()
def main():
args = parse_args()
cfg = load_config()
if args.problem:
cfg["problem"] = args.problem
if args.gpu is not None:
cfg["inference"]["gpu"] = args.gpu
if args.device:
cfg["inference"]["device"] = args.device
if args.n_col_domain is not None:
cfg["inference"]["n_col_domain"] = args.n_col_domain
if args.n_train_per_bc is not None:
cfg["data"]["n_train_per_bc"] = args.n_train_per_bc
if args.checkpoint_path:
cfg["inference"]["checkpoint_path"] = args.checkpoint_path
if args.output_dir:
cfg["inference"]["output_dir"] = args.output_dir
os.chdir(PROJECT_ROOT)
device = select_device(cfg["inference"])
models, metadata = build_models(
cfg,
device,
n_col_domain=cfg["inference"].get("n_col_domain", cfg["data"]["n_col_domain"]),
n_train_per_bc=cfg["data"]["n_train_per_bc"],
problem=cfg["problem"],
)
checkpoint = load_checkpoint(cfg["inference"]["checkpoint_path"], models, device)
share_mean_module(models)
for model in models:
model.eval()
clear_cached_kernels(models)
with torch.no_grad():
fields = predict_fields(models)
output_dir = resolve_path(cfg["inference"]["output_dir"])
output_dir.mkdir(parents=True, exist_ok=True)
npz_path = output_dir / "predictions.npz"
np.savez(npz_path, **tensor_to_numpy_dict(fields))
summary = {
**metadata,
"checkpoint_metadata": checkpoint.get("metadata", {}),
"output_file": str(npz_path),
"field_shapes": {key: list(value.shape) for key, value in fields.items()},
}
(output_dir / "inference_summary.json").write_text(json.dumps(summary, indent=2), encoding="utf-8")
print(f"Saved predictions to {npz_path}")
print(json.dumps(summary["field_shapes"], indent=2))
if __name__ == "__main__":
main()
|