import os from pathlib import Path from types import SimpleNamespace import yaml PROJECT_ROOT = Path(__file__).resolve().parents[1] torch = None get_data = None DistributedManager = None replace_function = None get_model = None DerivLoss = None L2Loss = None def load_config(): with open(PROJECT_ROOT / "config" / "config.yaml", "r", encoding="utf-8") as f: return yaml.safe_load(f) def load_onescience(): global torch global get_data global DistributedManager global replace_function global get_model global DerivLoss global L2Loss if torch is not None: return import torch as torch_module from onescience.datapipes.cfd_benchmark.data_factory import get_data as get_data_fn from onescience.distributed.manager import DistributedManager as DistributedManager_cls from onescience.memory.checkpoint import replace_function as replace_function_fn from onescience.models.cfd_benchmark.model_factory import get_model as get_model_fn from onescience.utils.cfd_benchmark.loss import DerivLoss as DerivLoss_cls from onescience.utils.cfd_benchmark.loss import L2Loss as L2Loss_cls torch = torch_module get_data = get_data_fn DistributedManager = DistributedManager_cls replace_function = replace_function_fn get_model = get_model_fn DerivLoss = DerivLoss_cls L2Loss = L2Loss_cls def build_args(cfg): data = cfg["data"] model = cfg["model"] train = cfg["train"] return SimpleNamespace( lr=train["lr"], epochs=train["epochs"], weight_decay=train["weight_decay"], pct_start=train["pct_start"], batch_size=data["batch_size"], gpu=0, max_grad_norm=train["max_grad_norm"], derivloss=train["derivloss"], optimizer=train["optimizer"], scheduler=train["scheduler"], step_size=100, gamma=0.5, find_unused_parameters=False, use_checkpoint=train["use_checkpoint"], checkpoint_layers=train["checkpoint_layers"], resume=False, data_path=cfg["paths"]["data_dir"], loader=data["loader"], config_name="config_name", train_ratio=0.8, ntrain=data["ntrain"], ntest=data["ntest"], normalize=data["normalize"], norm_type=data["norm_type"], geotype=data["geotype"], time_input=False, space_dim=data["space_dim"], fun_dim=data["fun_dim"], out_dim=data["out_dim"], shapelist=None, downsamplex=data["downsamplex"], downsampley=data["downsampley"], downsamplez=1, radius=0.2, task=data["task"], T_in=10, T_out=10, model=model["name"], n_hidden=model["n_hidden"], n_layers=model["n_layers"], n_heads=model["n_heads"], act=model["act"], mlp_ratio=model["mlp_ratio"], dropout=model["dropout"], unified_pos=model["unified_pos"], ref=model["ref"], slice_num=model["slice_num"], modes=12, psi_dim=8, attn_type="nystrom", mwt_k=3, branch_depth=5, trunk_depth=6, hidden_channels=[], kernel_size=5, emb_dims=128, eval=0, save_name=train["save_name"], vis_num=cfg["inference"]["vis_num"], vis_bound=None, ) def get_device(args, dist): if dist.world_size == 1 and torch.cuda.is_available(): return torch.device(f"cuda:{args.gpu}") if dist.world_size > 1 and torch.cuda.is_available(): return dist.device return torch.device("cpu") def make_optimizer(args, model): if args.optimizer == "AdamW": return torch.optim.AdamW( model.parameters(), lr=args.lr, weight_decay=args.weight_decay ) if args.optimizer == "Adam": return torch.optim.Adam( model.parameters(), lr=args.lr, weight_decay=args.weight_decay ) raise ValueError("Optimizer only AdamW or Adam") def make_scheduler(args, optimizer, train_loader): if args.scheduler == "OneCycleLR": return torch.optim.lr_scheduler.OneCycleLR( optimizer, max_lr=args.lr, epochs=args.epochs, steps_per_epoch=len(train_loader), pct_start=args.pct_start, ) if args.scheduler == "CosineAnnealingLR": return torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.epochs) if args.scheduler == "StepLR": return torch.optim.lr_scheduler.StepLR( optimizer, step_size=args.step_size, gamma=args.gamma ) raise ValueError("Scheduler only OneCycleLR, CosineAnnealingLR or StepLR") def evaluate(model, test_loader, dataset, args, device): myloss = L2Loss(size_average=False) model.eval() rel_err = 0.0 with torch.no_grad(): for pos, fx, y in test_loader: x, fx, y = pos.to(device), fx.to(device), y.to(device) if args.fun_dim == 0: fx = None out = model(x, fx) if args.normalize: out = dataset.y_normalizer.decode(out) rel_err += myloss(out, y).item() return rel_err / args.ntest def train(args): load_onescience() DistributedManager.initialize() dist = DistributedManager() dataset, train_loader, test_loader, args.shapelist = get_data(args, dist) device = get_device(args, dist) model = get_model(args, device).to(device) if hasattr(dataset, "x_normalizer"): dataset.x_normalizer = dataset.x_normalizer.to(device) if hasattr(dataset, "y_normalizer"): dataset.y_normalizer = dataset.y_normalizer.to(device) optimizer = make_optimizer(args, model) scheduler = make_scheduler(args, optimizer, train_loader) checkpoint_dir = PROJECT_ROOT / "checkpoints" checkpoint_dir.mkdir(parents=True, exist_ok=True) checkpoint_path = checkpoint_dir / f"{args.save_name}.pt" checkpoint_layers = ( [layer.strip() for layer in args.checkpoint_layers.split(",") if layer.strip()] if args.use_checkpoint and args.checkpoint_layers else [] ) myloss = L2Loss(size_average=False) regloss = DerivLoss(size_average=False, shapelist=args.shapelist) if args.derivloss else None best_test_loss = float("inf") best_epoch = 0 print(args) print(model) print(f"Use device: {device}") for epoch in range(args.epochs): model.train() train_loss = 0.0 for pos, fx, y in train_loader: x, fx, y = pos.to(device), fx.to(device), y.to(device) if args.fun_dim == 0: fx = None with replace_function( module=model, replace_layers_list=checkpoint_layers, ddp_flag=(dist.world_size > 1), ): out = model(x, fx) if args.normalize: out = dataset.y_normalizer.decode(out) y = dataset.y_normalizer.decode(y) loss = myloss(out, y) if regloss is not None: loss = loss + 0.1 * regloss(out, y) train_loss += loss.item() optimizer.zero_grad() loss.backward() if args.max_grad_norm is not None: torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm) optimizer.step() if args.scheduler == "OneCycleLR": scheduler.step() if args.scheduler in {"CosineAnnealingLR", "StepLR"}: scheduler.step() train_loss = train_loss / args.ntrain rel_err = evaluate(model, test_loader, dataset, args, device) if rel_err < best_test_loss: best_test_loss = rel_err best_epoch = epoch torch.save( { "epoch": epoch, "model_state": model.state_dict(), "optimizer_state": optimizer.state_dict(), "scheduler_state": scheduler.state_dict(), "best_test_loss": best_test_loss, "best_epoch": best_epoch, "args": vars(args), }, checkpoint_path, ) if epoch % 10 == 0: print("Epoch {} Train loss : {:.5f}".format(epoch, train_loss)) print("rel_err:{}".format(rel_err)) print( "Training completed. Best model saved at epoch {} with rel_err: {:.5f}".format( best_epoch, best_test_loss ) ) def main(): os.chdir(PROJECT_ROOT) cfg = load_config() if cfg.get("runtime", {}).get("device", "cpu") == "cpu": os.environ["CUDA_VISIBLE_DEVICES"] = "" train(build_args(cfg)) if __name__ == "__main__": main()