scripts/train.py

import ctypes
import logging
import os
import random
import sys
import sysconfig
import time
from pathlib import Path


def preload_python_shared_library():
    """Make libpython visible to native extensions loaded with ctypes."""
    libdir = sysconfig.get_config_var("LIBDIR")
    version = sysconfig.get_config_var("VERSION")
    if not libdir or not version:
        return

    candidates = [
        Path(libdir) / f"libpython{version}.so.1.0",
        Path(libdir) / f"libpython{version}.so",
    ]
    for libpython in candidates:
        if libpython.exists():
            ctypes.CDLL(str(libpython), mode=ctypes.RTLD_GLOBAL)
            return


preload_python_shared_library()

import numpy as np
import torch
import torch.nn as nn
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel
from tqdm import tqdm

# 获取项目根目录（train.py上级的上级）
root_path = Path(__file__).parent.parent
sys.path.insert(0, str(root_path))

from model.graphViT import GraphViT
from onescience.distributed.manager import DistributedManager
from onescience.utils.YParams import YParams
from onescience.datapipes.cfd import EagleDatapipe



def save_best_model(model, optimizer, checkpoint_dir: str):
    Path(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(),
        },
        Path(checkpoint_dir) / "best_model.pth",
    )


def load_best_model(model, checkpoint_dir: str, device: torch.device):
    ckpt_path = Path(checkpoint_dir) / "best_model.pth"
    checkpoint = torch.load(ckpt_path, map_location=device, weights_only=False)
    state_dict = checkpoint.get("model_state_dict", checkpoint)
    model.load_state_dict(state_dict)


def setup_logging(rank: int):
    logging.basicConfig(
        level=logging.INFO if rank == 0 else logging.WARNING,
        format="%(asctime)s - %(levelname)s - %(message)s",
        stream=sys.stdout,
        force=True,
    )
    return logging.getLogger("train")


def get_loss(velocity, pressure, output, state_hat, target, mask, alpha):
    velocity = velocity[:, 1:]
    pressure = pressure[:, 1:]
    velocity_hat = state_hat[:, 1:, :, :2]
    pressure_hat = state_hat[:, 1:, :, 2:]
    mask = mask[:, 1:].unsqueeze(-1)

    loss_velocity = torch.sqrt(((velocity * mask - velocity_hat * mask) ** 2).mean(dim=-1)).mean()
    loss_pressure = torch.sqrt(((pressure * mask - pressure_hat * mask) ** 2).mean(dim=-1)).mean()
    mse = nn.MSELoss()
    loss = mse(target[..., :2] * mask, output[..., :2] * mask)
    loss = loss + alpha * mse(target[..., 2:] * mask, output[..., 2:] * mask)
    return {"loss": loss, "MSE_velocity": loss_velocity, "MSE_pressure": loss_pressure}


def move_batch(x, device):
    return {
        "mesh_pos": x["mesh_pos"].to(device),
        "edges": x["edges"].to(device).long(),
        "velocity": x["velocity"].to(device),
        "pressure": x["pressure"].to(device),
        "node_type": x["node_type"].to(device),
        "mask": x["mask"].to(device),
        "cluster": x["cluster"].to(device).long(),
        "cluster_mask": x["cluster_mask"].to(device).long(),
    }


def fix_single_cluster_path(datapipe, cfg_data):
    if int(cfg_data.data.n_cluster) != 1:
        return

    cluster_path = Path(cfg_data.source.cluster_dir)
    for dataset_name in ("train_dataset", "val_dataset", "test_dataset"):
        dataset = getattr(datapipe, dataset_name, None)
        if dataset is not None and getattr(dataset, "cluster_path", None) is None:
            dataset.cluster_path = cluster_path


def validate(model, dataloader, device, alpha, manager):
    model.eval()
    total_loss, count = 0.0, 0
    with torch.no_grad():
        for x in dataloader:
            if not x:
                continue
            batch = move_batch(x, device)
            state = torch.cat([batch["velocity"], batch["pressure"]], dim=-1)
            state_hat, output, target = model(
                batch["mesh_pos"],
                batch["edges"],
                state,
                batch["node_type"],
                batch["cluster"],
                batch["cluster_mask"],
                apply_noise=False,
            )
            dataset = dataloader.dataset
            state_hat[..., :2], state_hat[..., 2:] = dataset.denormalize(
                state_hat[..., :2], state_hat[..., 2:]
            )
            velocity, pressure = dataset.denormalize(batch["velocity"], batch["pressure"])
            costs = get_loss(velocity, pressure, output, state_hat, target, batch["mask"], alpha)
            if manager.world_size > 1:
                dist.all_reduce(costs["loss"], op=dist.ReduceOp.AVG)
            total_loss += costs["loss"].item()
            count += 1
    return total_loss / max(count, 1)


def main():
    os.chdir(root_path)
    DistributedManager.initialize()
    manager = DistributedManager()
    logger = setup_logging(manager.rank)

    config_path = root_path / "config" / "config.yaml"
    cfg_model = YParams(config_path, "model")
    cfg_data = YParams(config_path, "datapipe")
    cfg_train = YParams(config_path, "training")

    seed = int(cfg_train.get("seed", 0))
    torch.manual_seed(seed)
    random.seed(seed)
    np.random.seed(seed)

    datapipe = EagleDatapipe(params=cfg_data, distributed=(manager.world_size > 1))
    fix_single_cluster_path(datapipe, cfg_data)
    train_loader, train_sampler = datapipe.train_dataloader()
    val_loader, val_sampler = datapipe.val_dataloader()

    device_name = cfg_train.get("device", "auto")
    device = manager.device if device_name == "auto" else torch.device(device_name)
    model = GraphViT(state_size=cfg_model.state_size, w_size=cfg_model.w_size).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))
    best_valid_loss = float("inf")
    best_epoch = 0

    for epoch in range(int(cfg_train.max_epoch)):
        start = time.time()
        if manager.world_size > 1:
            train_sampler.set_epoch(epoch)
            if val_sampler:
                val_sampler.set_epoch(epoch)

        model.train()
        train_loss, count = 0.0, 0
        pbar = tqdm(train_loader, desc=f"Epoch {epoch + 1}", disable=(manager.rank != 0))
        for x in pbar:
            if not x:
                continue
            batch = move_batch(x, device)
            state = torch.cat([batch["velocity"], batch["pressure"]], dim=-1)
            state_hat, output, target = model(
                batch["mesh_pos"],
                batch["edges"],
                state,
                batch["node_type"],
                batch["cluster"],
                batch["cluster_mask"],
                apply_noise=True,
            )
            state_hat[..., :2], state_hat[..., 2:] = train_loader.dataset.denormalize(
                state_hat[..., :2], state_hat[..., 2:]
            )
            velocity, pressure = train_loader.dataset.denormalize(batch["velocity"], batch["pressure"])
            costs = get_loss(velocity, pressure, output, state_hat, target, batch["mask"], cfg_train.loss_alpha)

            optimizer.zero_grad()
            costs["loss"].backward()
            optimizer.step()

            train_loss += costs["loss"].item()
            count += 1
            pbar.set_postfix(loss=f"{costs['loss'].item():.6f}")

        train_loss /= max(count, 1)
        valid_loss = validate(model, val_loader, device, cfg_train.loss_alpha, manager)
        if manager.rank == 0:
            logger.info(
                "Epoch %s/%s | %.2fs | train %.6f | valid %.6f",
                epoch + 1,
                cfg_train.max_epoch,
                time.time() - start,
                train_loss,
                valid_loss,
            )
            if valid_loss < best_valid_loss:
                best_valid_loss = valid_loss
                best_epoch = epoch
                save_best_model(model, optimizer, cfg_train.checkpoint_dir)
                logger.info("Saved checkpoint to %s/best_model.pth", cfg_train.checkpoint_dir)
            if epoch - best_epoch > int(cfg_train.patience):
                break

    if manager.rank == 0:
        final_model = GraphViT(state_size=cfg_model.state_size, w_size=cfg_model.w_size).to(device)
        load_best_model(final_model, cfg_train.checkpoint_dir, device)
        test_loader, _ = datapipe.test_dataloader()
        test_loss = validate(final_model, test_loader, device, cfg_train.loss_alpha, manager)
        logger.info("Final test loss: %.6f", test_loss)

    manager.cleanup()


if __name__ == "__main__":
    main()