| from __future__ import annotations |
|
|
| import argparse |
| import csv |
| import json |
| import math |
| import random |
| import sys |
| from pathlib import Path |
| from typing import Any |
|
|
| import numpy as np |
| import yaml |
|
|
| try: |
| import torch |
| from torch import nn |
| from torch.utils.data import DataLoader, Dataset |
| except ImportError as exc: |
| raise SystemExit(f"PyTorch is required to run training: {exc}") from exc |
|
|
| PROJECT_ROOT = Path(__file__).resolve().parents[1] |
| if str(PROJECT_ROOT) not in sys.path: |
| sys.path.insert(0, str(PROJECT_ROOT)) |
|
|
| from model.FlowBenchFNO import build_model_from_config |
|
|
|
|
| class FlowBenchTensorDataset(Dataset): |
| def __init__( |
| self, |
| x: np.ndarray, |
| y: np.ndarray, |
| indices: np.ndarray, |
| x_mean: np.ndarray | None = None, |
| x_std: np.ndarray | None = None, |
| y_mean: np.ndarray | None = None, |
| y_std: np.ndarray | None = None, |
| normalize: bool = True, |
| ) -> None: |
| self.x = x |
| self.y = y |
| self.indices = indices.astype(np.int64) |
| self.normalize = normalize |
| self.x_mean = x_mean |
| self.x_std = x_std |
| self.y_mean = y_mean |
| self.y_std = y_std |
|
|
| def __len__(self) -> int: |
| return int(self.indices.shape[0]) |
|
|
| def __getitem__(self, item: int) -> tuple[torch.Tensor, torch.Tensor]: |
| idx = self.indices[item] |
| x = self.x[idx] |
| y = self.y[idx] |
| if self.normalize: |
| x = (x - self.x_mean) / self.x_std |
| y = (y - self.y_mean) / self.y_std |
| return torch.from_numpy(x.astype(np.float32)), torch.from_numpy(y.astype(np.float32)) |
|
|
|
|
| def load_config(path: str | Path) -> dict[str, Any]: |
| with open(path, "r", encoding="utf-8") as handle: |
| return yaml.safe_load(handle) |
|
|
|
|
| def resolve_path(path: str | Path, root: Path = PROJECT_ROOT) -> Path: |
| path = Path(path) |
| return path if path.is_absolute() else root / path |
|
|
|
|
| def set_seed(seed: int) -> None: |
| random.seed(seed) |
| np.random.seed(seed) |
| torch.manual_seed(seed) |
| if torch.cuda.is_available(): |
| torch.cuda.manual_seed_all(seed) |
|
|
|
|
| def select_device(name: str) -> torch.device: |
| if name == "auto": |
| return torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| return torch.device(name) |
|
|
|
|
| def load_npz_array(path: Path, key: str) -> np.ndarray: |
| if not path.exists(): |
| raise FileNotFoundError(f"Missing data file: {path}") |
| archive = np.load(path) |
| actual_key = key if key in archive.files else archive.files[0] |
| array = archive[actual_key] |
| if array.ndim != 4: |
| raise ValueError(f"Expected [N,C,H,W] array in {path}, got {array.shape}") |
| return array.astype(np.float32, copy=False) |
|
|
|
|
| def make_splits(n_samples: int, split_fraction: float, shuffle: bool, seed: int) -> tuple[np.ndarray, np.ndarray]: |
| indices = np.arange(n_samples, dtype=np.int64) |
| if shuffle: |
| rng = np.random.default_rng(seed) |
| rng.shuffle(indices) |
| n_train = int(math.floor(n_samples * split_fraction)) |
| return indices[:n_train], indices[n_train:] |
|
|
|
|
| def channel_stats(array: np.ndarray, indices: np.ndarray) -> tuple[np.ndarray, np.ndarray]: |
| subset = array[indices] |
| mean = subset.mean(axis=(0, 2, 3), keepdims=False).astype(np.float32)[:, None, None] |
| std = subset.std(axis=(0, 2, 3), keepdims=False).astype(np.float32)[:, None, None] |
| std = np.maximum(std, np.asarray(1.0e-6, dtype=np.float32)) |
| return mean, std |
|
|
|
|
| def denormalize(tensor: torch.Tensor, mean: torch.Tensor, std: torch.Tensor) -> torch.Tensor: |
| return tensor * std + mean |
|
|
|
|
| def compute_metrics(pred: torch.Tensor, target: torch.Tensor) -> dict[str, float]: |
| diff = pred - target |
| mse = torch.mean(diff.square()).item() |
| mae = torch.mean(torch.abs(diff)).item() |
| per_sample_num = torch.linalg.vector_norm(diff.flatten(1), dim=1) |
| per_sample_den = torch.linalg.vector_norm(target.flatten(1), dim=1).clamp_min(1.0e-12) |
| relative_l2 = torch.mean(per_sample_num / per_sample_den).item() |
| return {"mse": mse, "relative_l2": relative_l2, "mae": mae} |
|
|
|
|
| @torch.no_grad() |
| def evaluate( |
| model: nn.Module, |
| loader: DataLoader, |
| device: torch.device, |
| y_mean: torch.Tensor, |
| y_std: torch.Tensor, |
| max_batches: int | None = None, |
| ) -> dict[str, float]: |
| model.eval() |
| totals = {"mse": 0.0, "relative_l2": 0.0, "mae": 0.0} |
| count = 0 |
| for batch_idx, (x, y) in enumerate(loader): |
| if max_batches is not None and batch_idx >= max_batches: |
| break |
| x = x.to(device) |
| y = y.to(device) |
| pred = model(x) |
| pred_phys = denormalize(pred, y_mean, y_std) |
| y_phys = denormalize(y, y_mean, y_std) |
| metrics = compute_metrics(pred_phys, y_phys) |
| batch_size = x.shape[0] |
| for key, value in metrics.items(): |
| totals[key] += value * batch_size |
| count += batch_size |
| if count == 0: |
| return {key: float("nan") for key in totals} |
| return {key: value / count for key, value in totals.items()} |
|
|
|
|
| def save_json(path: Path, payload: dict[str, Any]) -> None: |
| path.parent.mkdir(parents=True, exist_ok=True) |
| with open(path, "w", encoding="utf-8") as handle: |
| json.dump(payload, handle, indent=2) |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser(description="Train FlowBench FNO baseline.") |
| parser.add_argument("--config", default=str(PROJECT_ROOT / "config/config.yaml")) |
| parser.add_argument("--max-epochs", type=int, default=None, help="Override config training.num_epochs.") |
| parser.add_argument("--sample-limit", type=int, default=None, help="Use the first N samples before splitting.") |
| parser.add_argument("--device", default=None, help="Override experiment.device.") |
| parser.add_argument("--checkpoint", default=None, help="Override best checkpoint path.") |
| args = parser.parse_args() |
|
|
| config = load_config(args.config) |
| seed = int(config["experiment"].get("seed", 42)) |
| set_seed(seed) |
| device = select_device(args.device or config["experiment"].get("device", "auto")) |
|
|
| data_cfg = config["data"] |
| data_root = resolve_path(data_cfg["root_dir"]) |
| x = load_npz_array(data_root / data_cfg["x_file"], data_cfg.get("x_key", "data")) |
| y = load_npz_array(data_root / data_cfg["y_file"], data_cfg.get("y_key", "data")) |
| if x.shape[0] != y.shape[0]: |
| raise ValueError(f"X/Y sample counts differ: {x.shape[0]} vs {y.shape[0]}") |
| if args.sample_limit is not None: |
| x = x[: args.sample_limit] |
| y = y[: args.sample_limit] |
|
|
| model_cfg = config["model"] |
| if x.shape[1] != int(model_cfg["in_channels"]): |
| raise ValueError(f"Config in_channels={model_cfg['in_channels']} but X has {x.shape[1]} channels") |
| if y.shape[1] != int(model_cfg["out_channels"]): |
| raise ValueError(f"Config out_channels={model_cfg['out_channels']} but Y has {y.shape[1]} channels") |
|
|
| train_idx, val_idx = make_splits( |
| x.shape[0], |
| float(data_cfg.get("split_fraction", 0.8)), |
| bool(data_cfg.get("shuffle", True)), |
| seed, |
| ) |
| x_mean, x_std = channel_stats(x, train_idx) |
| y_mean, y_std = channel_stats(y, train_idx) |
| normalize = bool(data_cfg.get("normalize", True)) |
|
|
| train_ds = FlowBenchTensorDataset(x, y, train_idx, x_mean, x_std, y_mean, y_std, normalize) |
| val_ds = FlowBenchTensorDataset(x, y, val_idx, x_mean, x_std, y_mean, y_std, normalize) |
| train_loader = DataLoader( |
| train_ds, |
| batch_size=int(config["training"]["batch_size"]), |
| shuffle=True, |
| num_workers=int(data_cfg.get("num_workers", 0)), |
| pin_memory=bool(data_cfg.get("pin_memory", True)) and device.type == "cuda", |
| ) |
| val_loader = DataLoader( |
| val_ds, |
| batch_size=int(config["training"]["batch_size"]), |
| shuffle=False, |
| num_workers=int(data_cfg.get("num_workers", 0)), |
| pin_memory=bool(data_cfg.get("pin_memory", True)) and device.type == "cuda", |
| ) |
|
|
| model = build_model_from_config(config).to(device) |
| optimizer = torch.optim.Adam( |
| model.parameters(), |
| lr=float(config["training"]["learning_rate"]), |
| weight_decay=float(config["training"].get("weight_decay", 0.0)), |
| ) |
| criterion = nn.MSELoss() |
| y_mean_t = torch.from_numpy(y_mean).to(device) |
| y_std_t = torch.from_numpy(y_std).to(device) |
| x_mean_t = torch.from_numpy(x_mean) |
| x_std_t = torch.from_numpy(x_std) |
|
|
| num_epochs = args.max_epochs or int(config["training"]["num_epochs"]) |
| best_path = resolve_path(args.checkpoint or config["training"]["checkpoint_path"]) |
| last_path = resolve_path(config["training"]["last_checkpoint_path"]) |
| metrics_path = resolve_path(config["training"]["metrics_path"]) |
| history_path = resolve_path(config["training"]["history_path"]) |
| best_path.parent.mkdir(parents=True, exist_ok=True) |
| last_path.parent.mkdir(parents=True, exist_ok=True) |
| history_path.parent.mkdir(parents=True, exist_ok=True) |
|
|
| best_mse = float("inf") |
| with open(history_path, "w", encoding="utf-8", newline="") as handle: |
| writer = csv.DictWriter(handle, fieldnames=["epoch", "train_loss", "val_mse", "val_relative_l2", "val_mae"]) |
| writer.writeheader() |
| for epoch in range(1, num_epochs + 1): |
| model.train() |
| running_loss = 0.0 |
| seen = 0 |
| for x_batch, y_batch in train_loader: |
| x_batch = x_batch.to(device) |
| y_batch = y_batch.to(device) |
| optimizer.zero_grad(set_to_none=True) |
| pred = model(x_batch) |
| loss = criterion(pred, y_batch) |
| loss.backward() |
| clip_norm = config["training"].get("grad_clip_norm") |
| if clip_norm is not None: |
| torch.nn.utils.clip_grad_norm_(model.parameters(), float(clip_norm)) |
| optimizer.step() |
| batch_size = x_batch.shape[0] |
| running_loss += loss.item() * batch_size |
| seen += batch_size |
| train_loss = running_loss / max(1, seen) |
| val_metrics = evaluate(model, val_loader, device, y_mean_t, y_std_t) |
| row = { |
| "epoch": epoch, |
| "train_loss": train_loss, |
| "val_mse": val_metrics["mse"], |
| "val_relative_l2": val_metrics["relative_l2"], |
| "val_mae": val_metrics["mae"], |
| } |
| writer.writerow(row) |
| handle.flush() |
| print( |
| f"epoch={epoch:04d} train_loss={train_loss:.6e} " |
| f"val_mse={val_metrics['mse']:.6e} " |
| f"val_relative_l2={val_metrics['relative_l2']:.6e} " |
| f"val_mae={val_metrics['mae']:.6e}", |
| flush=True, |
| ) |
|
|
| checkpoint = { |
| "epoch": epoch, |
| "model_state_dict": model.state_dict(), |
| "optimizer_state_dict": optimizer.state_dict(), |
| "best_val_mse": min(best_mse, val_metrics["mse"]), |
| "config": config, |
| "normalization": { |
| "x_mean": x_mean_t.numpy(), |
| "x_std": x_std_t.numpy(), |
| "y_mean": torch.from_numpy(y_mean), |
| "y_std": torch.from_numpy(y_std), |
| }, |
| "splits": {"train_idx": train_idx, "val_idx": val_idx}, |
| "metrics": {"train_loss": train_loss, **val_metrics}, |
| } |
| torch.save(checkpoint, last_path) |
| if val_metrics["mse"] < best_mse: |
| best_mse = val_metrics["mse"] |
| torch.save(checkpoint, best_path) |
|
|
| final_payload = { |
| "best_val_mse": best_mse, |
| "best_checkpoint": str(best_path), |
| "last_checkpoint": str(last_path), |
| "history": str(history_path), |
| "num_samples": int(x.shape[0]), |
| "train_samples": int(train_idx.shape[0]), |
| "val_samples": int(val_idx.shape[0]), |
| "device": str(device), |
| } |
| save_json(metrics_path, final_payload) |
| print(json.dumps(final_payload, indent=2), flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
