| from __future__ import annotations |
|
|
| import argparse |
| import csv |
| import json |
| import random |
| import sys |
| import time |
| from pathlib import Path |
| from typing import Any |
|
|
| import numpy as np |
| import torch |
| from torch import nn |
| from torch.utils.data import DataLoader, Dataset |
| import yaml |
|
|
| PROJECT_ROOT = Path(__file__).resolve().parents[1] |
| if str(PROJECT_ROOT) not in sys.path: |
| sys.path.insert(0, str(PROJECT_ROOT)) |
|
|
| from model.neural_convection_fno import build_model_from_config |
|
|
|
|
| class TensorPairDataset(Dataset): |
| def __init__(self, x: torch.Tensor, y: torch.Tensor, indices: np.ndarray, stats: dict[str, torch.Tensor]) -> None: |
| self.x = x |
| self.y = y |
| self.indices = indices.astype(np.int64) |
| self.stats = stats |
|
|
| def __len__(self) -> int: |
| return int(self.indices.shape[0]) |
|
|
| def __getitem__(self, item: int) -> tuple[torch.Tensor, torch.Tensor]: |
| idx = int(self.indices[item]) |
| return normalize_x(self.x[idx], self.stats), normalize_y(self.y[idx], self.stats) |
|
|
|
|
| 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: str | Path) -> Path: |
| path = Path(path) |
| return path if path.is_absolute() else PROJECT_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_pair(root: Path, split: str, x_file: str, y_file: str) -> tuple[torch.Tensor, torch.Tensor]: |
| x_path = root / split / x_file |
| y_path = root / split / y_file |
| if not x_path.exists() or not y_path.exists(): |
| raise FileNotFoundError(f"Missing tensors for split {split}: {x_path}, {y_path}") |
| x = torch.load(x_path, map_location="cpu") |
| y = torch.load(y_path, map_location="cpu") |
| if not isinstance(x, torch.Tensor) or not isinstance(y, torch.Tensor): |
| raise TypeError(f"Expected torch.Tensor objects in {x_path} and {y_path}") |
| x = x.float().contiguous() |
| y = y.float().contiguous() |
| if tuple(x.shape) != tuple(y.shape): |
| raise ValueError(f"X/Y shape mismatch for {split}: {tuple(x.shape)} vs {tuple(y.shape)}") |
| if x.ndim != 4 or x.shape[1:] != (4, 64, 64): |
| raise ValueError(f"Expected [N,4,64,64] tensors for {split}, got {tuple(x.shape)}") |
| return x, y |
|
|
|
|
| def make_train_val_indices(n: int, val_fraction: float, seed: int, shuffle: bool) -> tuple[np.ndarray, np.ndarray]: |
| indices = np.arange(n, dtype=np.int64) |
| if shuffle: |
| rng = np.random.default_rng(seed) |
| rng.shuffle(indices) |
| n_val = int(round(n * val_fraction)) |
| n_train = n - n_val |
| return indices[:n_train], indices[n_train:] |
|
|
|
|
| def channel_minmax(tensor: torch.Tensor, indices: np.ndarray) -> tuple[torch.Tensor, torch.Tensor]: |
| subset = tensor[torch.as_tensor(indices, dtype=torch.long)] |
| mins = subset.amin(dim=(0, 2, 3)).view(1, -1, 1, 1) |
| maxs = subset.amax(dim=(0, 2, 3)).view(1, -1, 1, 1) |
| return mins, maxs |
|
|
|
|
| def build_stats(x: torch.Tensor, y: torch.Tensor, train_idx: np.ndarray) -> dict[str, torch.Tensor]: |
| x_min, x_max = channel_minmax(x, train_idx) |
| y_min, y_max = channel_minmax(y, train_idx) |
| return { |
| "x_min": x_min, |
| "x_max": x_max, |
| "x_range": (x_max - x_min).clamp_min(1.0e-8), |
| "y_min": y_min, |
| "y_max": y_max, |
| "y_range": (y_max - y_min).clamp_min(1.0e-8), |
| "train_indices": torch.as_tensor(train_idx, dtype=torch.long), |
| } |
|
|
|
|
| def normalize_x(x: torch.Tensor, stats: dict[str, torch.Tensor]) -> torch.Tensor: |
| return (x - stats["x_min"].squeeze(0)) / stats["x_range"].squeeze(0) |
|
|
|
|
| def normalize_y(y: torch.Tensor, stats: dict[str, torch.Tensor]) -> torch.Tensor: |
| return (y - stats["y_min"].squeeze(0)) / stats["y_range"].squeeze(0) |
|
|
|
|
| def stats_to_device(stats: dict[str, torch.Tensor], device: torch.device) -> dict[str, torch.Tensor]: |
| return {key: value.to(device) if isinstance(value, torch.Tensor) else value for key, value in stats.items()} |
|
|
|
|
| def count_parameters(model: nn.Module) -> int: |
| return sum(param.numel() for param in model.parameters() if param.requires_grad) |
|
|
|
|
| def tensor_to_numpy(tensor: torch.Tensor, dtype: np.dtype = np.float32) -> np.ndarray: |
| """Convert tensors without relying on torch's NumPy ABI bridge.""" |
| return np.asarray(tensor.detach().cpu().tolist(), dtype=dtype) |
|
|
|
|
| def ssim_2d(pred: np.ndarray, target: np.ndarray) -> float: |
| try: |
| from skimage.metrics import structural_similarity |
|
|
| return float(structural_similarity(target, pred, data_range=1.0)) |
| except Exception: |
| mu_x = float(pred.mean()) |
| mu_y = float(target.mean()) |
| sigma_x = float(pred.var()) |
| sigma_y = float(target.var()) |
| sigma_xy = float(((pred - mu_x) * (target - mu_y)).mean()) |
| c1 = 0.01**2 |
| c2 = 0.03**2 |
| return float(((2 * mu_x * mu_y + c1) * (2 * sigma_xy + c2)) / ((mu_x**2 + mu_y**2 + c1) * (sigma_x + sigma_y + c2))) |
|
|
|
|
| @torch.no_grad() |
| def evaluate_loader( |
| model: nn.Module, |
| loader: DataLoader, |
| device: torch.device, |
| channel_names: list[str], |
| collect: bool = False, |
| ) -> tuple[dict[str, Any], dict[str, np.ndarray] | None]: |
| model.eval() |
| total_loss = 0.0 |
| total_count = 0 |
| sqerr_sum = torch.zeros(len(channel_names), dtype=torch.float64) |
| elem_count = 0 |
| preds: list[np.ndarray] = [] |
| targets: list[np.ndarray] = [] |
| inputs: list[np.ndarray] = [] |
| start = time.perf_counter() |
| for x, y in loader: |
| x = x.to(device, non_blocking=True) |
| y = y.to(device, non_blocking=True) |
| pred = model(x) |
| loss = torch.mean((pred - y) ** 2) |
| batch = x.shape[0] |
| total_loss += float(loss.item()) * batch |
| total_count += batch |
| sqerr_sum += (pred.detach().cpu().double() - y.detach().cpu().double()).square().sum(dim=(0, 2, 3)) |
| elem_count += batch * y.shape[-1] * y.shape[-2] |
| if collect: |
| preds.append(tensor_to_numpy(pred)) |
| targets.append(tensor_to_numpy(y)) |
| inputs.append(tensor_to_numpy(x)) |
| elapsed = time.perf_counter() - start |
| mse = np.asarray((sqerr_sum / max(1, elem_count)).tolist(), dtype=np.float64) |
| metrics: dict[str, Any] = { |
| "loss": total_loss / max(1, total_count), |
| "per_channel": {}, |
| "time_seconds_per_sample": elapsed / max(1, total_count), |
| } |
| pred_all = target_all = input_all = None |
| if collect: |
| pred_all = np.concatenate(preds, axis=0) |
| target_all = np.concatenate(targets, axis=0) |
| input_all = np.concatenate(inputs, axis=0) |
| for idx, name in enumerate(channel_names): |
| ssim = None |
| if collect and pred_all is not None and target_all is not None: |
| ssim = float(np.mean([ssim_2d(pred_all[i, idx], target_all[i, idx]) for i in range(pred_all.shape[0])])) |
| metrics["per_channel"][name] = {"mse": float(mse[idx]), "ssim": ssim} |
| arrays = None |
| if collect and pred_all is not None and target_all is not None and input_all is not None: |
| arrays = {"pred_norm": pred_all, "target_norm": target_all, "input_norm": input_all} |
| return metrics, arrays |
|
|
|
|
| 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 compact FNO for 2D NeuralConvection.") |
| parser.add_argument("--config", default=str(PROJECT_ROOT / "config/config.yaml")) |
| parser.add_argument("--epochs", type=int, default=None) |
| parser.add_argument("--device", default=None) |
| args = parser.parse_args() |
|
|
| config = load_config(args.config) |
| seed = int(config["experiment"].get("seed", 2026)) |
| set_seed(seed) |
| device = select_device(args.device or config["experiment"].get("device", "auto")) |
| print(f"device={device}", flush=True) |
|
|
| data_cfg = config["data"] |
| root = resolve(data_cfg["root_dir"]) |
| channel_names = list(data_cfg["channel_names"]) |
| x_train, y_train = load_pair(root, data_cfg["train_split"], data_cfg["x_file"], data_cfg["y_file"]) |
| train_idx, val_idx = make_train_val_indices( |
| x_train.shape[0], |
| float(data_cfg["validation_fraction"]), |
| seed, |
| bool(data_cfg.get("shuffle_train", True)), |
| ) |
| stats = build_stats(x_train, y_train, train_idx) |
| stats_path = resolve(data_cfg["stats_path"]) |
| stats_path.parent.mkdir(parents=True, exist_ok=True) |
| torch.save({**stats, "val_indices": torch.as_tensor(val_idx, dtype=torch.long)}, stats_path) |
| print( |
| f"loaded train tensors X/Y={tuple(x_train.shape)}; train={len(train_idx)} val={len(val_idx)}; stats={stats_path}", |
| flush=True, |
| ) |
|
|
| train_ds = TensorPairDataset(x_train, y_train, train_idx, stats) |
| val_ds = TensorPairDataset(x_train, y_train, val_idx, stats) |
| 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) |
| print(f"model_parameters={count_parameters(model)}", flush=True) |
| optimizer = torch.optim.AdamW( |
| model.parameters(), |
| lr=float(config["training"]["learning_rate"]), |
| weight_decay=float(config["training"]["weight_decay"]), |
| ) |
| scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( |
| optimizer, |
| mode=config["training"]["scheduler"].get("mode", "min"), |
| factor=float(config["training"]["scheduler"].get("factor", 0.5)), |
| patience=int(config["training"]["scheduler"].get("patience", 20)), |
| ) |
| criterion = nn.MSELoss() |
| epochs = int(args.epochs or config["training"]["epochs"]) |
| grad_clip = float(config["training"]["grad_clip_norm"]) |
| best_val = float("inf") |
| best_epoch = 0 |
| ckpt_path = resolve(config["training"]["checkpoint_path"]) |
| last_path = resolve(config["training"]["last_checkpoint_path"]) |
| ckpt_path.parent.mkdir(parents=True, exist_ok=True) |
| history_path = resolve(config["training"]["history_path"]) |
| history_path.parent.mkdir(parents=True, exist_ok=True) |
|
|
| with open(history_path, "w", encoding="utf-8", newline="") as handle: |
| writer = csv.DictWriter(handle, fieldnames=["epoch", "train_loss", "val_loss", "lr"]) |
| writer.writeheader() |
| for epoch in range(1, epochs + 1): |
| model.train() |
| train_total = 0.0 |
| train_count = 0 |
| for x, y in train_loader: |
| x = x.to(device, non_blocking=True) |
| y = y.to(device, non_blocking=True) |
| optimizer.zero_grad(set_to_none=True) |
| pred = model(x) |
| loss = criterion(pred, y) |
| loss.backward() |
| torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip) |
| optimizer.step() |
| batch = x.shape[0] |
| train_total += float(loss.item()) * batch |
| train_count += batch |
| train_loss = train_total / max(1, train_count) |
| val_metrics, _ = evaluate_loader(model, val_loader, device, channel_names, collect=False) |
| val_loss = float(val_metrics["loss"]) |
| scheduler.step(val_loss) |
| lr = float(optimizer.param_groups[0]["lr"]) |
| writer.writerow({"epoch": epoch, "train_loss": train_loss, "val_loss": val_loss, "lr": lr}) |
| handle.flush() |
| print(f"epoch={epoch:03d} train_loss={train_loss:.8e} val_loss={val_loss:.8e} lr={lr:.6e}", flush=True) |
| payload = { |
| "epoch": epoch, |
| "model_state_dict": model.state_dict(), |
| "optimizer_state_dict": optimizer.state_dict(), |
| "config": config, |
| "best_val_loss": min(best_val, val_loss), |
| "normalization_stats_path": str(stats_path), |
| } |
| torch.save(payload, last_path) |
| if val_loss < best_val: |
| best_val = val_loss |
| best_epoch = epoch |
| torch.save(payload, ckpt_path) |
| print(f"saved_best_checkpoint={ckpt_path} best_val_loss={best_val:.8e}", flush=True) |
|
|
| checkpoint = torch.load(ckpt_path, map_location=device) |
| model.load_state_dict(checkpoint["model_state_dict"]) |
| stats_device = stats_to_device(stats, device) |
| del stats_device |
|
|
| metrics_out: dict[str, Any] = { |
| "best_epoch": best_epoch, |
| "best_val_loss": best_val, |
| "checkpoint_path": str(ckpt_path), |
| "stats_path": str(stats_path), |
| "splits": {"train": int(len(train_idx)), "val": int(len(val_idx))}, |
| "tests": {}, |
| "paper_table5_compact_fno_2d": config["paper_table5_compact_fno_2d"], |
| } |
| for label, split in data_cfg["test_splits"].items(): |
| x_test, y_test = load_pair(root, split, data_cfg["x_file"], data_cfg["y_file"]) |
| test_idx = np.arange(x_test.shape[0], dtype=np.int64) |
| test_ds = TensorPairDataset(x_test, y_test, test_idx, stats) |
| test_loader = DataLoader( |
| test_ds, |
| batch_size=int(config["inference"]["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", |
| ) |
| collect = label == "in_distribution" |
| test_metrics, arrays = evaluate_loader(model, test_loader, device, channel_names, collect=collect) |
| metrics_out["tests"][label] = test_metrics |
| print(f"test_split={label} loss={test_metrics['loss']:.8e} time_s_per_sample={test_metrics['time_seconds_per_sample']:.6e}", flush=True) |
| for channel, vals in test_metrics["per_channel"].items(): |
| print(f" {channel}: mse={vals['mse']:.8e} ssim={vals['ssim']}", flush=True) |
| if collect and arrays is not None: |
| np.savez_compressed(resolve(config["inference"]["output_path"]), **arrays, channel_names=np.array(channel_names)) |
| save_json(resolve(config["training"]["metrics_path"]), metrics_out) |
| print(f"metrics_path={resolve(config['training']['metrics_path'])}", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
