File size: 5,088 Bytes
338c3e4 | 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 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 | import argparse
import json
import importlib.util
import os
import sys
from pathlib import Path
import torch
from tqdm import tqdm
PROJECT_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(PROJECT_ROOT))
from model import build_model, infer_task_type
import onescience
from onescience.distributed.manager import DistributedManager
from onescience.utils.YParams import YParams
def resolve_path(path_value):
path = Path(path_value)
return path if path.is_absolute() else PROJECT_ROOT / path
def parse_args():
parser = argparse.ArgumentParser(description="Run CFDBench inference.")
parser.add_argument("--model", default=None, help="Override root.model.name and read weight/<model>.pt by default.")
return parser.parse_args()
def load_config(model_name=None):
cfg = YParams(str(PROJECT_ROOT / "conf" / "config.yaml"), "root")
if model_name:
cfg.model.name = model_name
elif os.environ.get("CFDBENCH_MODEL_NAME"):
cfg.model.name = os.environ["CFDBENCH_MODEL_NAME"]
cfg.datapipe.source.data_dir = str(resolve_path(cfg.datapipe.source.data_dir))
cfg.inference.output_dir = str(resolve_path(cfg.inference.output_dir))
return cfg
def checkpoint_path(cfg):
path_value = cfg.inference.get("checkpoint_path", "auto")
if path_value == "auto":
return PROJECT_ROOT / "weight" / f"{cfg.model.name}.pt"
return resolve_path(path_value)
def output_path(cfg, task_type):
output_dir = Path(cfg.inference.output_dir)
if cfg.inference.get("group_by_model", False):
output_dir = output_dir / task_type / cfg.model.name
return output_dir
def select_device(requested, dist):
if requested == "auto":
return dist.device
if requested.startswith("cuda") and not torch.cuda.is_available():
raise RuntimeError(f"Requested device {requested!r}, but CUDA is not available.")
return torch.device(requested)
def load_cfdbench_datapipe_class():
runtime_root = Path(onescience.__file__).resolve().parent
datapipe_file = runtime_root / "datapipes" / "cfd" / "cfdbench.py"
spec = importlib.util.spec_from_file_location("_onescience_cfdbench_datapipe", datapipe_file)
if spec is None or spec.loader is None:
raise ImportError(f"Cannot load CFDBench datapipe from {datapipe_file}")
module = importlib.util.module_from_spec(spec)
sys.modules[spec.name] = module
spec.loader.exec_module(module)
return module.CFDBenchDatapipe
def predict_full_frame(model, batch, task_type):
if task_type == "auto":
return model.generate(
inputs=batch["inputs"],
case_params=batch["case_params"],
mask=batch.get("mask"),
)
labels = batch["label"]
return model.generate_one(
case_params=batch["case_params"],
t=batch["t"],
height=labels.shape[-2],
width=labels.shape[-1],
)
def main():
args = parse_args()
DistributedManager.initialize()
dist = DistributedManager()
cfg = load_config(args.model)
task_type = infer_task_type(cfg.model.name)
cfg.datapipe.data.task_type = task_type
device = select_device(cfg.inference.get("device", "auto"), dist)
ckpt_path = checkpoint_path(cfg)
if not ckpt_path.is_file():
raise FileNotFoundError(f"Missing checkpoint: {ckpt_path}. Run scripts/train.py first or update root.inference.checkpoint_path.")
torch.set_num_threads(1)
CFDBenchDatapipe = load_cfdbench_datapipe_class()
datapipe = CFDBenchDatapipe(cfg.datapipe, distributed=False)
model = build_model(cfg).to(device)
model.load_state_dict(torch.load(ckpt_path, map_location=device, weights_only=True))
model.eval()
output_dir = output_path(cfg, task_type)
output_dir.mkdir(parents=True, exist_ok=True)
preds = []
score_lists = {name: [] for name in model.loss_fn.get_score_names()}
with torch.no_grad():
for step, batch in enumerate(tqdm(datapipe.test_dataloader(), desc="Inferring")):
if step >= cfg.inference.max_batches:
break
batch = {key: value.to(device) for key, value in batch.items()}
pred = predict_full_frame(model, batch, task_type)
label = batch["label"][:, : pred.shape[1]]
loss = model.loss_fn(preds=pred, labels=label)
preds.append(pred.detach().cpu())
for key, value in loss.items():
score_lists[key].append(float(value.detach().cpu()))
if not preds:
raise RuntimeError("Inference loader is empty. Check dataset split and data directory.")
torch.save(torch.cat(preds, dim=0), output_dir / "preds.pt")
means = {key: sum(values) / len(values) for key, values in score_lists.items() if values}
with open(output_dir / "scores.json", "w", encoding="utf-8") as f:
json.dump({"mean": means, "all": score_lists}, f, indent=2)
print(f"Saved predictions to {output_dir / 'preds.pt'}")
print(f"Scores: {means}")
DistributedManager.cleanup()
if __name__ == "__main__":
main()
|