scripts/train.py

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())