from __future__ import annotations import logging import random import sys import time from pathlib import Path import numpy as np import torch import torch.distributed as dist import torch.nn as nn from torch.nn.parallel import DistributedDataParallel from torch.utils.data.distributed import DistributedSampler from torch_geometric.loader import DataLoader as PyGDataLoader ROOT = Path(__file__).resolve().parents[1] sys.path.insert(0, str(ROOT)) from model import Transolver2D, Transolver2D_plus from onescience.distributed.manager import DistributedManager from onescience.utils.YParams import YParams def resolve_path(path: str) -> Path: value = Path(path).expanduser() return value if value.is_absolute() else ROOT / value def setup_logging(rank: int) -> logging.Logger: logging.basicConfig( level=logging.INFO if rank == 0 else logging.WARNING, format="%(asctime)s - %(levelname)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S", ) return logging.getLogger("transolver-train") def load_graph(path: Path): try: return torch.load(path, weights_only=False) except TypeError: return torch.load(path) class FakeAirfRANSDataset(torch.utils.data.Dataset): def __init__(self, data_dir: Path, stats_dir: Path, split: str, coef_norm=None): manifest_path = data_dir / "manifest.json" if not manifest_path.exists(): raise FileNotFoundError(f"Fake manifest not found: {manifest_path}. Run scripts/fake_data.py first.") import json manifest = json.loads(manifest_path.read_text(encoding="utf-8")) self.data_dir = data_dir self.names = manifest[split] self.coef_norm = coef_norm or ( np.load(stats_dir / "mean_in.npy"), np.load(stats_dir / "std_in.npy"), np.load(stats_dir / "mean_out.npy"), np.load(stats_dir / "std_out.npy"), ) def __len__(self) -> int: return len(self.names) def __getitem__(self, index: int): data = load_graph(self.data_dir / f"{self.names[index]}.pt") mean_in, std_in, mean_out, std_out = self.coef_norm data.x = (data.x - torch.as_tensor(mean_in, dtype=data.x.dtype)) / torch.as_tensor(std_in + 1e-8, dtype=data.x.dtype) data.y = (data.y - torch.as_tensor(mean_out, dtype=data.y.dtype)) / torch.as_tensor(std_out + 1e-8, dtype=data.y.dtype) return data class FakeAirfRANSDatapipe: def __init__(self, cfg_data, distributed: bool): self.cfg_data = cfg_data self.distributed = distributed data_dir = resolve_path(cfg_data.source.data_dir) stats_dir = resolve_path(cfg_data.source.stats_dir) self.train_dataset = FakeAirfRANSDataset(data_dir, stats_dir, cfg_data.data.splits.train_name) self.coef_norm = self.train_dataset.coef_norm self.val_dataset = FakeAirfRANSDataset(data_dir, stats_dir, cfg_data.data.splits.val_name, self.coef_norm) self.test_dataset = FakeAirfRANSDataset(data_dir, stats_dir, cfg_data.data.splits.test_name, self.coef_norm) def _loader(self, dataset, shuffle: bool): sampler = DistributedSampler(dataset, shuffle=shuffle) if self.distributed else None return ( PyGDataLoader( dataset, batch_size=self.cfg_data.dataloader.batch_size, num_workers=self.cfg_data.dataloader.num_workers, pin_memory=torch.cuda.is_available(), shuffle=shuffle and sampler is None, sampler=sampler, ), sampler, ) def train_dataloader(self): return self._loader(self.train_dataset, True) def val_dataloader(self): return self._loader(self.val_dataset, False) def build_model(cfg_model): model_name = cfg_model.name if model_name not in cfg_model.specific_params: raise ValueError(f"Model '{model_name}' is missing from model.specific_params") params = cfg_model.specific_params[model_name] model_cls = Transolver2D_plus if model_name == "Transolver_plus" else Transolver2D if model_name not in ("Transolver", "Transolver_plus"): raise NotImplementedError("This refactored package localizes only Transolver and Transolver_plus.") return model_cls( n_hidden=params.n_hidden, n_layers=params.n_layers, space_dim=params.space_dim, fun_dim=params.fun_dim, n_head=params.n_head, mlp_ratio=params.mlp_ratio, out_dim=params.out_dim, slice_num=params.slice_num, unified_pos=bool(params.unified_pos), ) def build_datapipe(cfg_data, model_params, distributed: bool): cfg_data.model_hparams = model_params if cfg_data.backend == "airfrans": from onescience.datapipes.cfd import AirfRANSDatapipe return AirfRANSDatapipe(params=cfg_data, distributed=distributed) if cfg_data.backend == "fake_airfrans": return FakeAirfRANSDatapipe(cfg_data, distributed=distributed) raise ValueError(f"Unsupported datapipe.backend: {cfg_data.backend}") def masked_mean(losses: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: if mask.any(): return losses[mask].mean() return losses.mean() * 0.0 def save_checkpoint(model, optimizer, scheduler, epoch: int, loss: float, checkpoint_dir: Path, model_name: str) -> None: checkpoint_dir.mkdir(parents=True, exist_ok=True) model_to_save = model.module if hasattr(model, "module") else model torch.save( { "model_state_dict": model_to_save.state_dict(), "optimizer_state_dict": optimizer.state_dict(), "scheduler_state_dict": scheduler.state_dict(), "epoch": epoch, "loss": loss, }, checkpoint_dir / f"{model_name}.pth", ) def main() -> int: config_path = ROOT / "conf" / "config.yaml" cfg_model = YParams(str(config_path), "model") cfg_data = YParams(str(config_path), "datapipe") cfg_train = YParams(str(config_path), "training") torch.manual_seed(int(cfg_train.seed)) np.random.seed(int(cfg_train.seed)) random.seed(int(cfg_train.seed)) DistributedManager.initialize() manager = DistributedManager() logger = setup_logging(manager.rank) device_name = getattr(cfg_train, "device", "auto") if device_name == "cpu": device = torch.device("cpu") elif device_name.startswith("cuda"): device = torch.device(device_name if torch.cuda.is_available() else "cpu") else: device = torch.device(f"cuda:{cfg_train.gpuid}" if torch.cuda.is_available() and manager.world_size == 1 else manager.device) model_params = cfg_model.specific_params[cfg_model.name] datapipe = build_datapipe(cfg_data, model_params, distributed=(manager.world_size > 1)) train_loader, train_sampler = datapipe.train_dataloader() val_loader, val_sampler = datapipe.val_dataloader() if len(train_loader) == 0 or len(val_loader) == 0: raise RuntimeError("Train and validation loaders must both be non-empty.") model = build_model(cfg_model).to(device) if manager.world_size > 1: model = DistributedDataParallel(model, device_ids=[manager.local_rank], output_device=manager.local_rank) optimizer = torch.optim.Adam(model.parameters(), lr=float(cfg_train.lr)) scheduler = torch.optim.lr_scheduler.OneCycleLR( optimizer, max_lr=float(cfg_train.lr), total_steps=max(1, len(train_loader) * int(cfg_train.max_epoch)), ) criterion = nn.L1Loss(reduction="none") if cfg_train.loss_criterion == "MAE" else nn.MSELoss(reduction="none") use_weighted_loss = cfg_train.loss_criterion == "MSE_weighted" checkpoint_dir = resolve_path(cfg_train.checkpoint_dir) best_valid_loss = float("inf") best_epoch = -1 if manager.rank == 0: params = sum(p.numel() for p in model.parameters() if p.requires_grad) logger.info("Training %s with %.3fM parameters on %s", cfg_model.name, params / 1e6, device) for epoch in range(int(cfg_train.max_epoch)): start = time.time() if train_sampler is not None: train_sampler.set_epoch(epoch) if val_sampler is not None: val_sampler.set_epoch(epoch) model.train() train_loss = 0.0 for data in train_loader: data = data.to(device) optimizer.zero_grad(set_to_none=True) out = model(data) losses = criterion(out, data.y) surf_loss = masked_mean(losses, data.surf) vol_loss = masked_mean(losses, ~data.surf) loss = vol_loss + float(cfg_train.loss_weight) * surf_loss if use_weighted_loss else losses.mean() loss.backward() optimizer.step() scheduler.step() train_loss += float(loss.detach().cpu()) train_loss /= len(train_loader) model.eval() valid_loss = 0.0 with torch.no_grad(): for data in val_loader: data = data.to(device) out = model(data) losses = criterion(out, data.y) surf_loss = masked_mean(losses, data.surf) vol_loss = masked_mean(losses, ~data.surf) loss = vol_loss + float(cfg_train.loss_weight) * surf_loss if use_weighted_loss else losses.mean() if manager.world_size > 1: dist.all_reduce(loss, op=dist.ReduceOp.AVG) valid_loss += float(loss.detach().cpu()) valid_loss /= len(val_loader) if manager.rank == 0: logger.info( "Epoch %d/%d train=%.6f valid=%.6f time=%.2fs", epoch + 1, int(cfg_train.max_epoch), train_loss, valid_loss, time.time() - start, ) if valid_loss < best_valid_loss: best_valid_loss = valid_loss best_epoch = epoch save_checkpoint(model, optimizer, scheduler, epoch, valid_loss, checkpoint_dir, cfg_model.name) logger.info("Saved checkpoint to %s", checkpoint_dir / f"{cfg_model.name}.pth") if epoch - best_epoch > int(cfg_train.patience): logger.info("Early stopping after %d epochs without improvement.", int(cfg_train.patience)) break if manager.world_size > 1 and dist.is_available() and dist.is_initialized(): dist.barrier() dist.destroy_process_group() return 0 if __name__ == "__main__": raise SystemExit(main())