OneScience-Group/FNO
07
1#!/usr/bin/env python32"""Train the paper-specified recurrent FNO-2D Navier--Stokes model."""3 4from __future__ import annotations5 6import argparse7import json8import os9import platform10import random11import sys12import time13from copy import deepcopy14from datetime import datetime, timezone15from pathlib import Path16from typing import Any17 18import numpy as np19import scipy20import scipy.io21import torch22import yaml23from torch import Tensor, nn24from torch.utils.data import DataLoader, TensorDataset25 26 27PROJECT_ROOT = Path(__file__).resolve().parents[1]28if str(PROJECT_ROOT) not in sys.path:29 sys.path.insert(0, str(PROJECT_ROOT))30 31from models import build_model_from_config # noqa: E40232 33 34def parse_args() -> argparse.Namespace:35 parser = argparse.ArgumentParser(36 description="Train FNO-2D on the validated Navier-Stokes trajectory MAT file."37 )38 parser.add_argument(39 "--config",40 type=Path,41 default=PROJECT_ROOT / "config" / "config.yaml",42 help="YAML configuration path.",43 )44 parser.add_argument("--epochs", type=int, default=None, help="Override epochs (smoke only).")45 parser.add_argument(46 "--rollout-steps", type=int, default=None, help="Override rollout steps (smoke only)."47 )48 parser.add_argument(49 "--max-train-samples", type=int, default=None, help="Limit train samples (smoke only)."50 )51 parser.add_argument(52 "--max-test-samples", type=int, default=None, help="Limit test samples (smoke only)."53 )54 parser.add_argument("--batch-size", type=int, default=None, help="Override batch size.")55 parser.add_argument(56 "--device", default="auto", help="auto, cpu, cuda, or an explicit torch device."57 )58 parser.add_argument("--checkpoint", type=Path, default=None, help="Output checkpoint path.")59 parser.add_argument(60 "--resume",61 type=Path,62 default=None,63 help="Resume from a formal training-state checkpoint.",64 )65 parser.add_argument("--output-dir", type=Path, default=None, help="Output directory.")66 parser.add_argument(67 "--run-type", choices=("formal", "smoke"), default="formal", help="Run provenance tag."68 )69 return parser.parse_args()70 71 72def load_config(path: Path) -> dict[str, Any]:73 config_path = path.expanduser().resolve()74 if not config_path.is_file():75 raise FileNotFoundError(f"Configuration file does not exist: {config_path}")76 with config_path.open("r", encoding="utf-8") as handle:77 config = yaml.safe_load(handle)78 if not isinstance(config, dict):79 raise ValueError("Configuration root must be a mapping")80 validate_config(config)81 return config82 83 84def validate_config(config: dict[str, Any]) -> None:85 for section in ("paper", "data", "model", "training", "paths"):86 if section not in config or not isinstance(config[section], dict):87 raise ValueError(f"Missing configuration section: {section}")88 89 data = config["data"]90 model = config["model"]91 training = config["training"]92 required_data = {93 "root",94 "file",95 "key",96 "layout",97 "dtype",98 "expected_shape",99 "resolution",100 "ntrain",101 "ntest",102 "test_start",103 "history",104 "horizon",105 "normalization",106 }107 required_model = {108 "input_channels",109 "output_channels",110 "width",111 "modes1",112 "modes2",113 "num_layers",114 "projection_width",115 "use_grid",116 }117 required_training = {118 "epochs",119 "batch_size",120 "optimizer",121 "learning_rate",122 "weight_decay",123 "scheduler",124 "scheduler_step_size",125 "scheduler_gamma",126 "seed",127 "dtype",128 "relative_l2_epsilon",129 "checkpoint_monitor",130 }131 for label, mapping, required in (132 ("data", data, required_data),133 ("model", model, required_model),134 ("training", training, required_training),135 ):136 missing = sorted(required.difference(mapping))137 if missing:138 raise ValueError(f"Missing {label} configuration keys: {missing}")139 140 if data["layout"] != "N,H,W,T":141 raise ValueError("This reproduction requires data.layout=N,H,W,T")142 if data["dtype"] != "float32" or training["dtype"] != "float32":143 raise ValueError("The audited reproduction requires float32 data and training")144 if data["normalization"] != "none":145 raise ValueError("The paper-faithful default requires normalization=none")146 if int(data["history"]) != 10 or int(model["input_channels"]) != 10:147 raise ValueError("The FNO-2D experiment requires ten history channels")148 if int(data["horizon"]) != 10 or int(model["output_channels"]) != 1:149 raise ValueError("The experiment requires a ten-step rollout of one-step outputs")150 if int(model["width"]) != 32 or int(model["num_layers"]) != 4:151 raise ValueError("Strict paper settings require width=32 and num_layers=4")152 if int(model["modes1"]) != 12 or int(model["modes2"]) != 12:153 raise ValueError("Strict paper settings require 12 retained modes per axis")154 if str(training["optimizer"]).lower() != "adam":155 raise ValueError("The paper specifies Adam")156 if str(training["scheduler"]).lower() != "step_lr":157 raise ValueError("The paper schedule is represented by StepLR")158 if str(training["checkpoint_monitor"]) != "train_full_relative_l2":159 raise ValueError("Test metrics must not select the checkpoint")160 161 162def resolve_project_path(value: str | Path) -> Path:163 path = Path(value).expanduser()164 return path.resolve() if path.is_absolute() else (PROJECT_ROOT / path).resolve()165 166 167def data_file_from_config(config: dict[str, Any]) -> Path:168 data = config["data"]169 path = Path(str(data["root"])).expanduser() / str(data["file"])170 path = path.resolve()171 if not path.is_file():172 raise FileNotFoundError(f"Navier-Stokes MAT file does not exist: {path}")173 return path174 175 176def load_trajectory_array(config: dict[str, Any]) -> np.ndarray:177 data = config["data"]178 path = data_file_from_config(config)179 key = str(data["key"])180 payload = scipy.io.loadmat(path, variable_names=[key])181 if key not in payload:182 raise KeyError(f"MAT field {key!r} is missing from {path}")183 trajectories = payload[key]184 expected_shape = tuple(int(value) for value in data["expected_shape"])185 if trajectories.shape != expected_shape:186 raise ValueError(187 f"Expected {key} shape {expected_shape}, received {trajectories.shape}"188 )189 if trajectories.dtype != np.float32:190 raise TypeError(f"Expected {key} dtype float32, received {trajectories.dtype}")191 if not np.isfinite(trajectories).all():192 raise ValueError("Trajectory array contains NaN or Inf")193 return trajectories194 195 196def build_datasets(197 config: dict[str, Any],198 max_train_samples: int | None = None,199 max_test_samples: int | None = None,200 rollout_steps: int | None = None,201) -> tuple[TensorDataset, TensorDataset]:202 data = config["data"]203 trajectories = load_trajectory_array(config)204 history = int(data["history"])205 horizon = int(data["horizon"])206 steps = horizon if rollout_steps is None else int(rollout_steps)207 if not 1 <= steps <= horizon:208 raise ValueError(f"rollout_steps must be in [1,{horizon}], got {steps}")209 210 ntrain = int(data["ntrain"])211 ntest = int(data["ntest"])212 train_start = int(data.get("train_start", 0))213 test_start = int(data["test_start"])214 train_count = ntrain if max_train_samples is None else min(ntrain, max_train_samples)215 test_count = ntest if max_test_samples is None else min(ntest, max_test_samples)216 if train_count <= 0 or test_count <= 0:217 raise ValueError("Training and test sample counts must be positive")218 219 train = trajectories[train_start : train_start + train_count]220 test = trajectories[test_start : test_start + test_count]221 train_history = torch.from_numpy(train[..., :history])222 train_target = torch.from_numpy(train[..., history : history + steps])223 test_history = torch.from_numpy(test[..., :history])224 test_target = torch.from_numpy(test[..., history : history + steps])225 expected_hw = tuple(int(value) for value in data["resolution"])226 expected_history_shape = (expected_hw[0], expected_hw[1], history)227 if tuple(train_history.shape[1:]) != expected_history_shape:228 raise ValueError(f"Invalid train history shape: {train_history.shape}")229 if tuple(test_history.shape[1:]) != expected_history_shape:230 raise ValueError(f"Invalid test history shape: {test_history.shape}")231 return TensorDataset(train_history, train_target), TensorDataset(test_history, test_target)232 233 234def seed_everything(seed: int, deterministic: bool) -> None:235 random.seed(seed)236 np.random.seed(seed)237 torch.manual_seed(seed)238 if torch.cuda.is_available():239 torch.cuda.manual_seed_all(seed)240 if hasattr(torch.backends, "cudnn"):241 torch.backends.cudnn.deterministic = deterministic242 torch.backends.cudnn.benchmark = not deterministic243 if deterministic:244 torch.use_deterministic_algorithms(True, warn_only=True)245 246 247def select_device(requested: str) -> torch.device:248 if requested == "auto":249 return torch.device("cuda" if torch.cuda.is_available() else "cpu")250 device = torch.device(requested)251 if device.type == "cuda" and not torch.cuda.is_available():252 raise RuntimeError(f"Requested {requested}, but torch reports no CUDA/DCU device")253 return device254 255 256def relative_l2_per_sample(prediction: Tensor, target: Tensor, epsilon: float) -> Tensor:257 if prediction.shape != target.shape:258 raise ValueError(f"Relative L2 shape mismatch: {prediction.shape} vs {target.shape}")259 if prediction.ndim < 2:260 raise ValueError("Relative L2 inputs must include a batch and at least one feature axis")261 difference = torch.linalg.vector_norm(262 (prediction - target).reshape(prediction.shape[0], -1), dim=1263 )264 denominator = torch.linalg.vector_norm(target.reshape(target.shape[0], -1), dim=1)265 return difference / (denominator + epsilon)266 267 268def autoregressive_rollout(269 model: nn.Module,270 history: Tensor,271 target: Tensor,272 epsilon: float,273) -> tuple[Tensor, Tensor, Tensor]:274 if history.ndim != 4 or target.ndim != 4:275 raise ValueError("history and target must be channel-last four-dimensional tensors")276 window = history277 predictions: list[Tensor] = []278 step_ratios: list[Tensor] = []279 for step in range(target.shape[-1]):280 prediction = model(window)281 expected_step_shape = (*history.shape[:-1], 1)282 if tuple(prediction.shape) != expected_step_shape:283 raise ValueError(284 f"Model returned {tuple(prediction.shape)}, expected {expected_step_shape}"285 )286 target_step = target[..., step : step + 1]287 predictions.append(prediction)288 step_ratios.append(relative_l2_per_sample(prediction, target_step, epsilon))289 window = torch.cat((window[..., 1:], prediction), dim=-1)290 rollout = torch.cat(predictions, dim=-1)291 ratios = torch.stack(step_ratios, dim=1)292 backward_loss = ratios.mean(dim=0).sum()293 return rollout, ratios, backward_loss294 295 296def build_loader(297 dataset: TensorDataset,298 batch_size: int,299 shuffle: bool,300 num_workers: int,301 pin_memory: bool,302 seed: int,303) -> DataLoader:304 generator = torch.Generator()305 generator.manual_seed(seed)306 return DataLoader(307 dataset,308 batch_size=batch_size,309 shuffle=shuffle,310 num_workers=num_workers,311 pin_memory=pin_memory,312 drop_last=False,313 generator=generator,314 )315 316 317def run_epoch(318 model: nn.Module,319 loader: DataLoader,320 device: torch.device,321 epsilon: float,322 optimizer: torch.optim.Optimizer | None,323) -> dict[str, float]:324 training = optimizer is not None325 model.train(training)326 total_samples = 0327 total_step_ratio = 0.0328 total_full_ratio = 0.0329 steps_per_sample: int | None = None330 331 context = torch.enable_grad() if training else torch.inference_mode()332 with context:333 for history, target in loader:334 history = history.to(device=device, dtype=torch.float32, non_blocking=True)335 target = target.to(device=device, dtype=torch.float32, non_blocking=True)336 if training:337 optimizer.zero_grad(set_to_none=True)338 prediction, step_ratios, backward_loss = autoregressive_rollout(339 model, history, target, epsilon340 )341 full_ratios = relative_l2_per_sample(prediction, target, epsilon)342 if not torch.isfinite(backward_loss) or not torch.isfinite(full_ratios).all():343 raise FloatingPointError("Non-finite training/evaluation loss encountered")344 if training:345 backward_loss.backward()346 for name, parameter in model.named_parameters():347 if parameter.grad is not None and not torch.isfinite(parameter.grad).all():348 raise FloatingPointError(f"Non-finite gradient in parameter {name}")349 optimizer.step()350 351 batch_samples = int(history.shape[0])352 total_samples += batch_samples353 total_step_ratio += float(step_ratios.detach().sum().cpu())354 total_full_ratio += float(full_ratios.detach().sum().cpu())355 steps_per_sample = int(step_ratios.shape[1])356 357 if total_samples == 0 or steps_per_sample is None:358 raise RuntimeError("DataLoader produced no batches")359 return {360 "step_loss_sum": total_step_ratio / total_samples,361 "mean_step_relative_l2": total_step_ratio / (total_samples * steps_per_sample),362 "full_relative_l2": total_full_ratio / total_samples,363 "samples": float(total_samples),364 "rollout_steps": float(steps_per_sample),365 }366 367 368def synchronize(device: torch.device) -> None:369 if device.type == "cuda" and torch.cuda.is_available():370 torch.cuda.synchronize(device)371 372 373def environment_metadata(device: torch.device) -> dict[str, Any]:374 device_name = "cpu"375 if device.type == "cuda" and torch.cuda.is_available():376 device_name = torch.cuda.get_device_name(device)377 return {378 "python": platform.python_version(),379 "pytorch": torch.__version__,380 "numpy": np.__version__,381 "scipy": scipy.__version__,382 "pyyaml": yaml.__version__,383 "device": str(device),384 "device_name": device_name,385 "torch_cuda_version": torch.version.cuda,386 "hostname": platform.node(),387 }388 389 390def atomic_write_json(path: Path, payload: dict[str, Any]) -> None:391 path.parent.mkdir(parents=True, exist_ok=True)392 temporary = path.with_suffix(path.suffix + ".tmp")393 with temporary.open("w", encoding="utf-8") as handle:394 json.dump(payload, handle, indent=2, ensure_ascii=False)395 handle.write("\n")396 os.replace(temporary, path)397 398 399def atomic_torch_save(path: Path, payload: dict[str, Any]) -> None:400 path.parent.mkdir(parents=True, exist_ok=True)401 temporary = path.with_suffix(path.suffix + ".tmp")402 torch.save(payload, temporary)403 os.replace(temporary, path)404 405 406def load_training_checkpoint(path: Path, device: torch.device) -> dict[str, Any]:407 checkpoint_path = path.expanduser().resolve()408 if not checkpoint_path.is_file():409 raise FileNotFoundError(f"Resume checkpoint does not exist: {checkpoint_path}")410 try:411 checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)412 except TypeError:413 checkpoint = torch.load(checkpoint_path, map_location=device)414 if not isinstance(checkpoint, dict):415 raise TypeError("Resume checkpoint root must be a mapping")416 required = {417 "epoch",418 "model_state_dict",419 "optimizer_state_dict",420 "scheduler_state_dict",421 "best_train_full_relative_l2",422 "monitor",423 "test_selected",424 "config",425 "run_type",426 }427 missing = sorted(required.difference(checkpoint))428 if missing:429 raise KeyError(f"Resume checkpoint is missing required keys: {missing}")430 if checkpoint["run_type"] != "formal":431 raise ValueError("Formal training can only resume a formal checkpoint")432 if checkpoint["monitor"] != "train_full_relative_l2":433 raise ValueError(f"Unexpected checkpoint monitor: {checkpoint['monitor']}")434 if checkpoint["test_selected"] is not False:435 raise ValueError("This reproduction forbids resuming a test-selected checkpoint")436 return checkpoint437 438 439def capture_rng_state(train_loader: DataLoader) -> dict[str, Any]:440 state: dict[str, Any] = {441 "python": random.getstate(),442 "numpy": np.random.get_state(),443 "torch": torch.get_rng_state(),444 "train_loader_generator": train_loader.generator.get_state(),445 }446 if torch.cuda.is_available():447 state["cuda"] = torch.cuda.get_rng_state_all()448 return state449 450 451def restore_rng_state(state: dict[str, Any], train_loader: DataLoader) -> None:452 random.setstate(state["python"])453 np.random.set_state(state["numpy"])454 torch.set_rng_state(state["torch"].cpu())455 train_loader.generator.set_state(state["train_loader_generator"].cpu())456 if torch.cuda.is_available() and "cuda" in state:457 torch.cuda.set_rng_state_all([value.cpu() for value in state["cuda"]])458 459 460def main() -> None:461 args = parse_args()462 config = load_config(args.config)463 training = config["training"]464 seed = int(training["seed"])465 seed_everything(seed, bool(training.get("deterministic", True)))466 device = select_device(args.device)467 468 epochs = int(training["epochs"] if args.epochs is None else args.epochs)469 batch_size = int(training["batch_size"] if args.batch_size is None else args.batch_size)470 horizon = int(config["data"]["horizon"])471 rollout_steps = horizon if args.rollout_steps is None else int(args.rollout_steps)472 if epochs <= 0 or batch_size <= 0:473 raise ValueError("epochs and batch_size must be positive")474 if args.run_type == "formal" and any(475 value is not None476 for value in (477 args.epochs,478 args.rollout_steps,479 args.max_train_samples,480 args.max_test_samples,481 args.output_dir,482 args.checkpoint,483 )484 ):485 raise ValueError(486 "Formal runs use the exact configured data, epochs, rollout, checkpoint, "487 "and results paths; overrides require --run-type smoke"488 )489 490 output_dir = (491 resolve_project_path(config["paths"]["results_dir"])492 if args.output_dir is None493 else args.output_dir.expanduser().resolve()494 )495 checkpoint_path = (496 resolve_project_path(config["paths"]["checkpoint"])497 if args.checkpoint is None498 else args.checkpoint.expanduser().resolve()499 )500 latest_checkpoint_path = checkpoint_path.with_name("last_model.pth")501 history_path = (502 resolve_project_path(config["paths"]["train_history"])503 if args.output_dir is None504 else output_dir / "train_history.json"505 )506 output_dir.mkdir(parents=True, exist_ok=True)507 checkpoint_path.parent.mkdir(parents=True, exist_ok=True)508 509 train_dataset, test_dataset = build_datasets(510 config,511 max_train_samples=args.max_train_samples,512 max_test_samples=args.max_test_samples,513 rollout_steps=rollout_steps,514 )515 pin_memory = bool(training.get("pin_memory", True)) and device.type == "cuda"516 train_loader = build_loader(517 train_dataset,518 batch_size,519 True,520 int(training.get("num_workers", 0)),521 pin_memory,522 seed,523 )524 test_loader = build_loader(525 test_dataset,526 batch_size,527 False,528 int(training.get("num_workers", 0)),529 pin_memory,530 seed,531 )532 533 # Move devices without forcing a global dtype conversion: the pointwise534 # parameters are float32 while spectral weights must remain complex64.535 model = build_model_from_config(config).to(device=device)536 parameter_count = sum(parameter.numel() for parameter in model.parameters())537 paper_parameter_count = int(config["paper"]["reference_parameter_count"])538 optimizer = torch.optim.Adam(539 model.parameters(),540 lr=float(training["learning_rate"]),541 weight_decay=float(training["weight_decay"]),542 )543 scheduler = torch.optim.lr_scheduler.StepLR(544 optimizer,545 step_size=int(training["scheduler_step_size"]),546 gamma=float(training["scheduler_gamma"]),547 )548 epsilon = float(training["relative_l2_epsilon"])549 550 run_started = datetime.now(timezone.utc).isoformat()551 history_payload: dict[str, Any] = {552 "schema_version": "fno-ns2d-train-history-v1",553 "run_type": args.run_type,554 "started_at": run_started,555 "completed_at": None,556 "config_path": str(args.config.expanduser().resolve()),557 "data_path": str(data_file_from_config(config)),558 "checkpoint_path": str(checkpoint_path),559 "config": deepcopy(config),560 "environment": environment_metadata(device),561 "parameter_count": parameter_count,562 "paper_parameter_count": paper_parameter_count,563 "parameter_count_difference": parameter_count - paper_parameter_count,564 "test_selected": False,565 "epochs_requested": epochs,566 "rollout_steps": rollout_steps,567 "train_samples": len(train_dataset),568 "test_samples": len(test_dataset),569 "history": [],570 }571 572 best_metric = float("inf")573 start_epoch = 1574 resume_exact_rng = True575 if args.resume is not None:576 resume_path = args.resume.expanduser().resolve()577 resume_checkpoint = load_training_checkpoint(resume_path, device)578 if resume_checkpoint["config"] != config:579 raise ValueError("Resume checkpoint configuration differs from the current YAML")580 resume_epoch = int(resume_checkpoint["epoch"])581 if not 1 <= resume_epoch < epochs:582 raise ValueError(583 f"Resume epoch must be in [1,{epochs - 1}], received {resume_epoch}"584 )585 if not history_path.is_file():586 raise FileNotFoundError(587 f"Training history required for audited resume is missing: {history_path}"588 )589 with history_path.open("r", encoding="utf-8") as handle:590 previous_history = json.load(handle)591 if not isinstance(previous_history, dict):592 raise TypeError("Existing training history root must be a mapping")593 records = previous_history.get("history")594 if not isinstance(records, list) or len(records) < resume_epoch:595 raise ValueError(596 f"Existing history has {len(records) if isinstance(records, list) else 0} "597 f"records, fewer than resume epoch {resume_epoch}"598 )599 retained_records = records[:resume_epoch]600 if int(retained_records[-1]["epoch"]) != resume_epoch:601 raise ValueError("Existing training history is not contiguous at the resume epoch")602 checkpoint_metric = float(resume_checkpoint["best_train_full_relative_l2"])603 history_best = min(float(record["train_full_relative_l2"]) for record in retained_records)604 if not np.isclose(checkpoint_metric, history_best, rtol=1e-7, atol=1e-9):605 raise ValueError(606 f"Resume checkpoint best metric {checkpoint_metric} differs from retained "607 f"history best {history_best}"608 )609 610 model.load_state_dict(resume_checkpoint["model_state_dict"], strict=True)611 optimizer.load_state_dict(resume_checkpoint["optimizer_state_dict"])612 scheduler.load_state_dict(resume_checkpoint["scheduler_state_dict"])613 if not bool(resume_checkpoint.get("scheduler_step_applied", False)):614 scheduler.step()615 rng_state = resume_checkpoint.get("rng_state")616 if isinstance(rng_state, dict):617 restore_rng_state(rng_state, train_loader)618 else:619 resume_exact_rng = False620 621 discarded_records = len(records) - resume_epoch622 history_payload = previous_history623 history_payload["history"] = retained_records624 history_payload["completed_at"] = None625 history_payload["completed"] = False626 history_payload["epochs_requested"] = epochs627 history_payload["resume_exact_rng"] = resume_exact_rng628 resume_events = history_payload.setdefault("resume_events", [])629 resume_events.append(630 {631 "resumed_at": run_started,632 "checkpoint_path": str(resume_path),633 "checkpoint_epoch": resume_epoch,634 "discarded_history_records": discarded_records,635 "exact_rng_state_restored": resume_exact_rng,636 }637 )638 best_metric = checkpoint_metric639 start_epoch = resume_epoch + 1640 atomic_write_json(history_path, history_payload)641 print(642 f"resume checkpoint={resume_path} epoch={resume_epoch} "643 f"discarded_history_records={discarded_records} "644 f"exact_rng_state_restored={'yes' if resume_exact_rng else 'no'}",645 flush=True,646 )647 print(648 f"device={device} parameters={parameter_count} "649 f"paper_parameters={paper_parameter_count} delta={parameter_count-paper_parameter_count}",650 flush=True,651 )652 653 if device.type == "cuda" and torch.cuda.is_available():654 torch.cuda.reset_peak_memory_stats(device)655 for epoch in range(start_epoch, epochs + 1):656 synchronize(device)657 epoch_start = time.perf_counter()658 lr_used = float(optimizer.param_groups[0]["lr"])659 train_metrics = run_epoch(model, train_loader, device, epsilon, optimizer)660 test_metrics = run_epoch(model, test_loader, device, epsilon, optimizer=None)661 synchronize(device)662 duration = time.perf_counter() - epoch_start663 peak_memory_bytes = (664 int(torch.cuda.max_memory_allocated(device))665 if device.type == "cuda" and torch.cuda.is_available()666 else 0667 )668 669 monitor = float(train_metrics["full_relative_l2"])670 is_best = monitor < best_metric671 if is_best:672 best_metric = monitor673 checkpoint = {674 "schema_version": "fno-ns2d-checkpoint-v1",675 "epoch": epoch,676 "model_state_dict": model.state_dict(),677 "optimizer_state_dict": optimizer.state_dict(),678 "scheduler_state_dict": scheduler.state_dict(),679 "best_train_full_relative_l2": best_metric,680 "test_full_relative_l2_at_best_epoch": float(681 test_metrics["full_relative_l2"]682 ),683 "monitor": "train_full_relative_l2",684 "test_selected": False,685 "config": deepcopy(config),686 "seed": seed,687 "parameter_count": parameter_count,688 "paper_parameter_count": paper_parameter_count,689 "parameter_count_difference": parameter_count - paper_parameter_count,690 "run_type": args.run_type,691 "scheduler_step_applied": False,692 "rng_state": capture_rng_state(train_loader),693 }694 atomic_torch_save(checkpoint_path, checkpoint)695 696 epoch_record = {697 "epoch": epoch,698 "learning_rate": lr_used,699 "duration_seconds": duration,700 "train_step_loss_sum": float(train_metrics["step_loss_sum"]),701 "train_mean_step_relative_l2": float(702 train_metrics["mean_step_relative_l2"]703 ),704 "train_full_relative_l2": float(train_metrics["full_relative_l2"]),705 "test_mean_step_relative_l2": float(test_metrics["mean_step_relative_l2"]),706 "test_full_relative_l2": float(test_metrics["full_relative_l2"]),707 "peak_accelerator_memory_bytes": peak_memory_bytes,708 "best": is_best,709 }710 history_payload["history"].append(epoch_record)711 history_payload["best_epoch"] = next(712 record["epoch"]713 for record in reversed(history_payload["history"])714 if record["best"]715 )716 history_payload["best_train_full_relative_l2"] = best_metric717 atomic_write_json(history_path, history_payload)718 print(719 f"epoch={epoch:04d}/{epochs:04d} time={duration:.3f}s lr={lr_used:.6g} "720 f"train_step_loss_sum={train_metrics['step_loss_sum']:.8f} "721 f"train_step_rel_l2={train_metrics['mean_step_relative_l2']:.8f} "722 f"train_full_rel_l2={train_metrics['full_relative_l2']:.8f} "723 f"test_step_rel_l2={test_metrics['mean_step_relative_l2']:.8f} "724 f"test_full_rel_l2={test_metrics['full_relative_l2']:.8f} "725 f"peak_mem_gib={peak_memory_bytes / (1024 ** 3):.3f} "726 f"best={'yes' if is_best else 'no'}",727 flush=True,728 )729 scheduler.step()730 latest_checkpoint = {731 "schema_version": "fno-ns2d-training-state-v1",732 "epoch": epoch,733 "model_state_dict": model.state_dict(),734 "optimizer_state_dict": optimizer.state_dict(),735 "scheduler_state_dict": scheduler.state_dict(),736 "best_train_full_relative_l2": best_metric,737 "latest_train_full_relative_l2": float(train_metrics["full_relative_l2"]),738 "latest_test_full_relative_l2": float(test_metrics["full_relative_l2"]),739 "monitor": "train_full_relative_l2",740 "test_selected": False,741 "config": deepcopy(config),742 "seed": seed,743 "parameter_count": parameter_count,744 "paper_parameter_count": paper_parameter_count,745 "parameter_count_difference": parameter_count - paper_parameter_count,746 "run_type": args.run_type,747 "scheduler_step_applied": True,748 "rng_state": capture_rng_state(train_loader),749 }750 atomic_torch_save(latest_checkpoint_path, latest_checkpoint)751 752 history_payload["completed_at"] = datetime.now(timezone.utc).isoformat()753 history_payload["completed"] = True754 atomic_write_json(history_path, history_payload)755 print(756 f"training_complete best_epoch={history_payload['best_epoch']} "757 f"best_train_full_rel_l2={best_metric:.8f} checkpoint={checkpoint_path}",758 flush=True,759 )760 761 762if __name__ == "__main__":763 main()764