import argparse import json import importlib.util import os import sys import time from pathlib import Path import numpy as np import torch from torch.nn.parallel import DistributedDataParallel as DDP from torch.optim import Adam, lr_scheduler 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="Train CFDBench static or autoregressive models.") parser.add_argument("--model", default=None, help="Override root.model.name, e.g. ffn, deeponet, fno, auto_ffn.") 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.training.output_dir = str(resolve_path(cfg.training.output_dir)) return cfg def checkpoint_path(cfg): name = cfg.training.get("checkpoint_name", "auto") if name == "auto": name = f"{cfg.model.name}.pt" return PROJECT_ROOT / "weight" / name def output_path(cfg, task_type): output_dir = Path(cfg.training.output_dir) if cfg.training.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 dump_json(data, path): path.parent.mkdir(parents=True, exist_ok=True) with open(path, "w", encoding="utf-8") as f: json.dump(data, f, indent=2) 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 mean_scores(score_lists): return {key: float(np.mean(values)) for key, values in score_lists.items() if values} def evaluate(model, loader, device, dist, desc="Evaluating"): model_eval = model.module if hasattr(model, "module") else model model_eval.eval() score_lists = {name: [] for name in model_eval.loss_fn.get_score_names()} with torch.no_grad(): iterator = tqdm(loader, desc=desc, disable=(dist.rank != 0)) for batch in iterator: batch = {key: value.to(device) for key, value in batch.items()} outputs = model_eval(**batch) for key, value in outputs["loss"].items(): score_lists[key].append(float(value.detach().cpu())) return {"mean": mean_scores(score_lists), "all": score_lists} def train(cfg, model, datapipe, output_dir, device, dist): train_cfg = cfg.training train_loader, train_sampler = datapipe.train_dataloader() val_loader, _ = datapipe.val_dataloader() if len(train_loader) == 0: raise RuntimeError("Training loader is empty. Increase fake_data.num_cases_per_subset or reduce batch_size.") optimizer = Adam(model.parameters(), lr=train_cfg.lr) scheduler = lr_scheduler.StepLR(optimizer, step_size=train_cfg.lr_step_size, gamma=train_cfg.lr_gamma) train_losses = [] best_nmse = float("inf") for epoch in range(train_cfg.num_epochs): if train_sampler: train_sampler.set_epoch(epoch) model.train() start = time.time() iterator = tqdm(train_loader, desc=f"Epoch {epoch}", disable=(dist.rank != 0)) for step, batch in enumerate(iterator): batch = {key: value.to(device) for key, value in batch.items()} outputs = model(**batch) loss = outputs["loss"][train_cfg.loss_name] optimizer.zero_grad(set_to_none=True) loss.backward() optimizer.step() loss_value = float(loss.detach().cpu()) train_losses.append(loss_value) if dist.rank == 0 and (step + 1) % train_cfg.log_interval == 0: iterator.set_postfix({"loss": f"{loss_value:.4e}"}) scheduler.step() if dist.rank == 0 and (epoch + 1) % train_cfg.eval_interval == 0: ckpt_dir = output_dir / f"ckpt-{epoch}" ckpt_dir.mkdir(parents=True, exist_ok=True) scores = evaluate(model, val_loader, device, dist, desc=f"Val {epoch}") dump_json(scores, ckpt_dir / "dev_scores.json") dump_json({"epoch": epoch, "train_loss": train_losses, "seconds": time.time() - start}, ckpt_dir / "scores.json") nmse = scores["mean"].get(train_cfg.loss_name, float("inf")) model_to_save = model.module if hasattr(model, "module") else model torch.save(model_to_save.state_dict(), ckpt_dir / "model.pt") if nmse <= best_nmse: best_nmse = nmse weight_path = checkpoint_path(cfg) torch.save(model_to_save.state_dict(), weight_path) print(f"Saved best checkpoint to {weight_path}") if dist.world_size > 1: torch.distributed.barrier() if dist.rank == 0: dump_json(train_losses, output_dir / "train_losses.json") def main(required_task_type="static", entry_name="scripts/train.py"): args = parse_args() DistributedManager.initialize() dist = DistributedManager() cfg = load_config(args.model) task_type = infer_task_type(cfg.model.name) if required_task_type is not None and task_type != required_task_type: expected = "ffn/deeponet" if required_task_type == "static" else "auto_* / resnet / unet / fno" raise ValueError( f"{entry_name} is the {required_task_type} entry, but model.name={cfg.model.name!r} is {task_type}. " f"Use a {expected} model, or run scripts/train_auto.py for autoregressive models." ) cfg.datapipe.data.task_type = task_type device = select_device(cfg.training.get("device", "auto"), dist) output_dir = output_path(cfg, task_type) if dist.rank == 0: output_dir.mkdir(parents=True, exist_ok=True) print(f"Config: {PROJECT_ROOT / 'conf' / 'config.yaml'}") print(f"Model: {cfg.model.name} ({task_type})") print(f"Data: {cfg.datapipe.source.data_dir}") print(f"Output: {output_dir}") print(f"Checkpoint: {checkpoint_path(cfg)}") torch.set_num_threads(1) CFDBenchDatapipe = load_cfdbench_datapipe_class() datapipe = CFDBenchDatapipe(cfg.datapipe, distributed=(dist.world_size > 1)) model = build_model(cfg).to(device) if dist.world_size > 1 and "train" in cfg.training.mode: device_ids = [dist.local_rank] if device.type == "cuda" else None model = DDP(model, device_ids=device_ids) if "train" in cfg.training.mode: train(cfg, model, datapipe, output_dir, device, dist) if "test" in cfg.training.mode and dist.rank == 0: scores = evaluate(model, datapipe.test_dataloader(), device, dist, desc="Test") dump_json(scores, output_dir / "test_scores.json") print(f"Test scores: {scores['mean']}") DistributedManager.cleanup() if __name__ == "__main__": main()