#!/usr/bin/env python3 """Train the paper-specified recurrent FNO-2D Navier--Stokes model.""" from __future__ import annotations import argparse import json import os import platform import random import sys import time from copy import deepcopy from datetime import datetime, timezone from pathlib import Path from typing import Any import numpy as np import scipy import scipy.io import torch import yaml from torch import Tensor, nn from torch.utils.data import DataLoader, TensorDataset PROJECT_ROOT = Path(__file__).resolve().parents[1] if str(PROJECT_ROOT) not in sys.path: sys.path.insert(0, str(PROJECT_ROOT)) from models import build_model_from_config # noqa: E402 def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description="Train FNO-2D on the validated Navier-Stokes trajectory MAT file." ) parser.add_argument( "--config", type=Path, default=PROJECT_ROOT / "config" / "config.yaml", help="YAML configuration path.", ) parser.add_argument("--epochs", type=int, default=None, help="Override epochs (smoke only).") parser.add_argument( "--rollout-steps", type=int, default=None, help="Override rollout steps (smoke only)." ) parser.add_argument( "--max-train-samples", type=int, default=None, help="Limit train samples (smoke only)." ) parser.add_argument( "--max-test-samples", type=int, default=None, help="Limit test samples (smoke only)." ) parser.add_argument("--batch-size", type=int, default=None, help="Override batch size.") parser.add_argument( "--device", default="auto", help="auto, cpu, cuda, or an explicit torch device." ) parser.add_argument("--checkpoint", type=Path, default=None, help="Output checkpoint path.") parser.add_argument( "--resume", type=Path, default=None, help="Resume from a formal training-state checkpoint.", ) parser.add_argument("--output-dir", type=Path, default=None, help="Output directory.") parser.add_argument( "--run-type", choices=("formal", "smoke"), default="formal", help="Run provenance tag." ) return parser.parse_args() def load_config(path: Path) -> dict[str, Any]: config_path = path.expanduser().resolve() if not config_path.is_file(): raise FileNotFoundError(f"Configuration file does not exist: {config_path}") with config_path.open("r", encoding="utf-8") as handle: config = yaml.safe_load(handle) if not isinstance(config, dict): raise ValueError("Configuration root must be a mapping") validate_config(config) return config def validate_config(config: dict[str, Any]) -> None: for section in ("paper", "data", "model", "training", "paths"): if section not in config or not isinstance(config[section], dict): raise ValueError(f"Missing configuration section: {section}") data = config["data"] model = config["model"] training = config["training"] required_data = { "root", "file", "key", "layout", "dtype", "expected_shape", "resolution", "ntrain", "ntest", "test_start", "history", "horizon", "normalization", } required_model = { "input_channels", "output_channels", "width", "modes1", "modes2", "num_layers", "projection_width", "use_grid", } required_training = { "epochs", "batch_size", "optimizer", "learning_rate", "weight_decay", "scheduler", "scheduler_step_size", "scheduler_gamma", "seed", "dtype", "relative_l2_epsilon", "checkpoint_monitor", } for label, mapping, required in ( ("data", data, required_data), ("model", model, required_model), ("training", training, required_training), ): missing = sorted(required.difference(mapping)) if missing: raise ValueError(f"Missing {label} configuration keys: {missing}") if data["layout"] != "N,H,W,T": raise ValueError("This reproduction requires data.layout=N,H,W,T") if data["dtype"] != "float32" or training["dtype"] != "float32": raise ValueError("The audited reproduction requires float32 data and training") if data["normalization"] != "none": raise ValueError("The paper-faithful default requires normalization=none") if int(data["history"]) != 10 or int(model["input_channels"]) != 10: raise ValueError("The FNO-2D experiment requires ten history channels") if int(data["horizon"]) != 10 or int(model["output_channels"]) != 1: raise ValueError("The experiment requires a ten-step rollout of one-step outputs") if int(model["width"]) != 32 or int(model["num_layers"]) != 4: raise ValueError("Strict paper settings require width=32 and num_layers=4") if int(model["modes1"]) != 12 or int(model["modes2"]) != 12: raise ValueError("Strict paper settings require 12 retained modes per axis") if str(training["optimizer"]).lower() != "adam": raise ValueError("The paper specifies Adam") if str(training["scheduler"]).lower() != "step_lr": raise ValueError("The paper schedule is represented by StepLR") if str(training["checkpoint_monitor"]) != "train_full_relative_l2": raise ValueError("Test metrics must not select the checkpoint") def resolve_project_path(value: str | Path) -> Path: path = Path(value).expanduser() return path.resolve() if path.is_absolute() else (PROJECT_ROOT / path).resolve() def data_file_from_config(config: dict[str, Any]) -> Path: data = config["data"] path = Path(str(data["root"])).expanduser() / str(data["file"]) path = path.resolve() if not path.is_file(): raise FileNotFoundError(f"Navier-Stokes MAT file does not exist: {path}") return path def load_trajectory_array(config: dict[str, Any]) -> np.ndarray: data = config["data"] path = data_file_from_config(config) key = str(data["key"]) payload = scipy.io.loadmat(path, variable_names=[key]) if key not in payload: raise KeyError(f"MAT field {key!r} is missing from {path}") trajectories = payload[key] expected_shape = tuple(int(value) for value in data["expected_shape"]) if trajectories.shape != expected_shape: raise ValueError( f"Expected {key} shape {expected_shape}, received {trajectories.shape}" ) if trajectories.dtype != np.float32: raise TypeError(f"Expected {key} dtype float32, received {trajectories.dtype}") if not np.isfinite(trajectories).all(): raise ValueError("Trajectory array contains NaN or Inf") return trajectories def build_datasets( config: dict[str, Any], max_train_samples: int | None = None, max_test_samples: int | None = None, rollout_steps: int | None = None, ) -> tuple[TensorDataset, TensorDataset]: data = config["data"] trajectories = load_trajectory_array(config) history = int(data["history"]) horizon = int(data["horizon"]) steps = horizon if rollout_steps is None else int(rollout_steps) if not 1 <= steps <= horizon: raise ValueError(f"rollout_steps must be in [1,{horizon}], got {steps}") ntrain = int(data["ntrain"]) ntest = int(data["ntest"]) train_start = int(data.get("train_start", 0)) test_start = int(data["test_start"]) train_count = ntrain if max_train_samples is None else min(ntrain, max_train_samples) test_count = ntest if max_test_samples is None else min(ntest, max_test_samples) if train_count <= 0 or test_count <= 0: raise ValueError("Training and test sample counts must be positive") train = trajectories[train_start : train_start + train_count] test = trajectories[test_start : test_start + test_count] train_history = torch.from_numpy(train[..., :history]) train_target = torch.from_numpy(train[..., history : history + steps]) test_history = torch.from_numpy(test[..., :history]) test_target = torch.from_numpy(test[..., history : history + steps]) expected_hw = tuple(int(value) for value in data["resolution"]) expected_history_shape = (expected_hw[0], expected_hw[1], history) if tuple(train_history.shape[1:]) != expected_history_shape: raise ValueError(f"Invalid train history shape: {train_history.shape}") if tuple(test_history.shape[1:]) != expected_history_shape: raise ValueError(f"Invalid test history shape: {test_history.shape}") return TensorDataset(train_history, train_target), TensorDataset(test_history, test_target) def seed_everything(seed: int, deterministic: bool) -> None: random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) if hasattr(torch.backends, "cudnn"): torch.backends.cudnn.deterministic = deterministic torch.backends.cudnn.benchmark = not deterministic if deterministic: torch.use_deterministic_algorithms(True, warn_only=True) def select_device(requested: str) -> torch.device: if requested == "auto": return torch.device("cuda" if torch.cuda.is_available() else "cpu") device = torch.device(requested) if device.type == "cuda" and not torch.cuda.is_available(): raise RuntimeError(f"Requested {requested}, but torch reports no CUDA/DCU device") return device def relative_l2_per_sample(prediction: Tensor, target: Tensor, epsilon: float) -> Tensor: if prediction.shape != target.shape: raise ValueError(f"Relative L2 shape mismatch: {prediction.shape} vs {target.shape}") if prediction.ndim < 2: raise ValueError("Relative L2 inputs must include a batch and at least one feature axis") difference = torch.linalg.vector_norm( (prediction - target).reshape(prediction.shape[0], -1), dim=1 ) denominator = torch.linalg.vector_norm(target.reshape(target.shape[0], -1), dim=1) return difference / (denominator + epsilon) def autoregressive_rollout( model: nn.Module, history: Tensor, target: Tensor, epsilon: float, ) -> tuple[Tensor, Tensor, Tensor]: if history.ndim != 4 or target.ndim != 4: raise ValueError("history and target must be channel-last four-dimensional tensors") window = history predictions: list[Tensor] = [] step_ratios: list[Tensor] = [] for step in range(target.shape[-1]): prediction = model(window) expected_step_shape = (*history.shape[:-1], 1) if tuple(prediction.shape) != expected_step_shape: raise ValueError( f"Model returned {tuple(prediction.shape)}, expected {expected_step_shape}" ) target_step = target[..., step : step + 1] predictions.append(prediction) step_ratios.append(relative_l2_per_sample(prediction, target_step, epsilon)) window = torch.cat((window[..., 1:], prediction), dim=-1) rollout = torch.cat(predictions, dim=-1) ratios = torch.stack(step_ratios, dim=1) backward_loss = ratios.mean(dim=0).sum() return rollout, ratios, backward_loss def build_loader( dataset: TensorDataset, batch_size: int, shuffle: bool, num_workers: int, pin_memory: bool, seed: int, ) -> DataLoader: generator = torch.Generator() generator.manual_seed(seed) return DataLoader( dataset, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers, pin_memory=pin_memory, drop_last=False, generator=generator, ) def run_epoch( model: nn.Module, loader: DataLoader, device: torch.device, epsilon: float, optimizer: torch.optim.Optimizer | None, ) -> dict[str, float]: training = optimizer is not None model.train(training) total_samples = 0 total_step_ratio = 0.0 total_full_ratio = 0.0 steps_per_sample: int | None = None context = torch.enable_grad() if training else torch.inference_mode() with context: for history, target in loader: history = history.to(device=device, dtype=torch.float32, non_blocking=True) target = target.to(device=device, dtype=torch.float32, non_blocking=True) if training: optimizer.zero_grad(set_to_none=True) prediction, step_ratios, backward_loss = autoregressive_rollout( model, history, target, epsilon ) full_ratios = relative_l2_per_sample(prediction, target, epsilon) if not torch.isfinite(backward_loss) or not torch.isfinite(full_ratios).all(): raise FloatingPointError("Non-finite training/evaluation loss encountered") if training: backward_loss.backward() for name, parameter in model.named_parameters(): if parameter.grad is not None and not torch.isfinite(parameter.grad).all(): raise FloatingPointError(f"Non-finite gradient in parameter {name}") optimizer.step() batch_samples = int(history.shape[0]) total_samples += batch_samples total_step_ratio += float(step_ratios.detach().sum().cpu()) total_full_ratio += float(full_ratios.detach().sum().cpu()) steps_per_sample = int(step_ratios.shape[1]) if total_samples == 0 or steps_per_sample is None: raise RuntimeError("DataLoader produced no batches") return { "step_loss_sum": total_step_ratio / total_samples, "mean_step_relative_l2": total_step_ratio / (total_samples * steps_per_sample), "full_relative_l2": total_full_ratio / total_samples, "samples": float(total_samples), "rollout_steps": float(steps_per_sample), } def synchronize(device: torch.device) -> None: if device.type == "cuda" and torch.cuda.is_available(): torch.cuda.synchronize(device) def environment_metadata(device: torch.device) -> dict[str, Any]: device_name = "cpu" if device.type == "cuda" and torch.cuda.is_available(): device_name = torch.cuda.get_device_name(device) return { "python": platform.python_version(), "pytorch": torch.__version__, "numpy": np.__version__, "scipy": scipy.__version__, "pyyaml": yaml.__version__, "device": str(device), "device_name": device_name, "torch_cuda_version": torch.version.cuda, "hostname": platform.node(), } def atomic_write_json(path: Path, payload: dict[str, Any]) -> None: path.parent.mkdir(parents=True, exist_ok=True) temporary = path.with_suffix(path.suffix + ".tmp") with temporary.open("w", encoding="utf-8") as handle: json.dump(payload, handle, indent=2, ensure_ascii=False) handle.write("\n") os.replace(temporary, path) def atomic_torch_save(path: Path, payload: dict[str, Any]) -> None: path.parent.mkdir(parents=True, exist_ok=True) temporary = path.with_suffix(path.suffix + ".tmp") torch.save(payload, temporary) os.replace(temporary, path) def load_training_checkpoint(path: Path, device: torch.device) -> dict[str, Any]: checkpoint_path = path.expanduser().resolve() if not checkpoint_path.is_file(): raise FileNotFoundError(f"Resume checkpoint does not exist: {checkpoint_path}") try: checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False) except TypeError: checkpoint = torch.load(checkpoint_path, map_location=device) if not isinstance(checkpoint, dict): raise TypeError("Resume checkpoint root must be a mapping") required = { "epoch", "model_state_dict", "optimizer_state_dict", "scheduler_state_dict", "best_train_full_relative_l2", "monitor", "test_selected", "config", "run_type", } missing = sorted(required.difference(checkpoint)) if missing: raise KeyError(f"Resume checkpoint is missing required keys: {missing}") if checkpoint["run_type"] != "formal": raise ValueError("Formal training can only resume a formal checkpoint") if checkpoint["monitor"] != "train_full_relative_l2": raise ValueError(f"Unexpected checkpoint monitor: {checkpoint['monitor']}") if checkpoint["test_selected"] is not False: raise ValueError("This reproduction forbids resuming a test-selected checkpoint") return checkpoint def capture_rng_state(train_loader: DataLoader) -> dict[str, Any]: state: dict[str, Any] = { "python": random.getstate(), "numpy": np.random.get_state(), "torch": torch.get_rng_state(), "train_loader_generator": train_loader.generator.get_state(), } if torch.cuda.is_available(): state["cuda"] = torch.cuda.get_rng_state_all() return state def restore_rng_state(state: dict[str, Any], train_loader: DataLoader) -> None: random.setstate(state["python"]) np.random.set_state(state["numpy"]) torch.set_rng_state(state["torch"].cpu()) train_loader.generator.set_state(state["train_loader_generator"].cpu()) if torch.cuda.is_available() and "cuda" in state: torch.cuda.set_rng_state_all([value.cpu() for value in state["cuda"]]) def main() -> None: args = parse_args() config = load_config(args.config) training = config["training"] seed = int(training["seed"]) seed_everything(seed, bool(training.get("deterministic", True))) device = select_device(args.device) epochs = int(training["epochs"] if args.epochs is None else args.epochs) batch_size = int(training["batch_size"] if args.batch_size is None else args.batch_size) horizon = int(config["data"]["horizon"]) rollout_steps = horizon if args.rollout_steps is None else int(args.rollout_steps) if epochs <= 0 or batch_size <= 0: raise ValueError("epochs and batch_size must be positive") if args.run_type == "formal" and any( value is not None for value in ( args.epochs, args.rollout_steps, args.max_train_samples, args.max_test_samples, args.output_dir, args.checkpoint, ) ): raise ValueError( "Formal runs use the exact configured data, epochs, rollout, checkpoint, " "and results paths; overrides require --run-type smoke" ) output_dir = ( resolve_project_path(config["paths"]["results_dir"]) if args.output_dir is None else args.output_dir.expanduser().resolve() ) checkpoint_path = ( resolve_project_path(config["paths"]["checkpoint"]) if args.checkpoint is None else args.checkpoint.expanduser().resolve() ) latest_checkpoint_path = checkpoint_path.with_name("last_model.pth") history_path = ( resolve_project_path(config["paths"]["train_history"]) if args.output_dir is None else output_dir / "train_history.json" ) output_dir.mkdir(parents=True, exist_ok=True) checkpoint_path.parent.mkdir(parents=True, exist_ok=True) train_dataset, test_dataset = build_datasets( config, max_train_samples=args.max_train_samples, max_test_samples=args.max_test_samples, rollout_steps=rollout_steps, ) pin_memory = bool(training.get("pin_memory", True)) and device.type == "cuda" train_loader = build_loader( train_dataset, batch_size, True, int(training.get("num_workers", 0)), pin_memory, seed, ) test_loader = build_loader( test_dataset, batch_size, False, int(training.get("num_workers", 0)), pin_memory, seed, ) # Move devices without forcing a global dtype conversion: the pointwise # parameters are float32 while spectral weights must remain complex64. model = build_model_from_config(config).to(device=device) parameter_count = sum(parameter.numel() for parameter in model.parameters()) paper_parameter_count = int(config["paper"]["reference_parameter_count"]) optimizer = torch.optim.Adam( model.parameters(), lr=float(training["learning_rate"]), weight_decay=float(training["weight_decay"]), ) scheduler = torch.optim.lr_scheduler.StepLR( optimizer, step_size=int(training["scheduler_step_size"]), gamma=float(training["scheduler_gamma"]), ) epsilon = float(training["relative_l2_epsilon"]) run_started = datetime.now(timezone.utc).isoformat() history_payload: dict[str, Any] = { "schema_version": "fno-ns2d-train-history-v1", "run_type": args.run_type, "started_at": run_started, "completed_at": None, "config_path": str(args.config.expanduser().resolve()), "data_path": str(data_file_from_config(config)), "checkpoint_path": str(checkpoint_path), "config": deepcopy(config), "environment": environment_metadata(device), "parameter_count": parameter_count, "paper_parameter_count": paper_parameter_count, "parameter_count_difference": parameter_count - paper_parameter_count, "test_selected": False, "epochs_requested": epochs, "rollout_steps": rollout_steps, "train_samples": len(train_dataset), "test_samples": len(test_dataset), "history": [], } best_metric = float("inf") start_epoch = 1 resume_exact_rng = True if args.resume is not None: resume_path = args.resume.expanduser().resolve() resume_checkpoint = load_training_checkpoint(resume_path, device) if resume_checkpoint["config"] != config: raise ValueError("Resume checkpoint configuration differs from the current YAML") resume_epoch = int(resume_checkpoint["epoch"]) if not 1 <= resume_epoch < epochs: raise ValueError( f"Resume epoch must be in [1,{epochs - 1}], received {resume_epoch}" ) if not history_path.is_file(): raise FileNotFoundError( f"Training history required for audited resume is missing: {history_path}" ) with history_path.open("r", encoding="utf-8") as handle: previous_history = json.load(handle) if not isinstance(previous_history, dict): raise TypeError("Existing training history root must be a mapping") records = previous_history.get("history") if not isinstance(records, list) or len(records) < resume_epoch: raise ValueError( f"Existing history has {len(records) if isinstance(records, list) else 0} " f"records, fewer than resume epoch {resume_epoch}" ) retained_records = records[:resume_epoch] if int(retained_records[-1]["epoch"]) != resume_epoch: raise ValueError("Existing training history is not contiguous at the resume epoch") checkpoint_metric = float(resume_checkpoint["best_train_full_relative_l2"]) history_best = min(float(record["train_full_relative_l2"]) for record in retained_records) if not np.isclose(checkpoint_metric, history_best, rtol=1e-7, atol=1e-9): raise ValueError( f"Resume checkpoint best metric {checkpoint_metric} differs from retained " f"history best {history_best}" ) model.load_state_dict(resume_checkpoint["model_state_dict"], strict=True) optimizer.load_state_dict(resume_checkpoint["optimizer_state_dict"]) scheduler.load_state_dict(resume_checkpoint["scheduler_state_dict"]) if not bool(resume_checkpoint.get("scheduler_step_applied", False)): scheduler.step() rng_state = resume_checkpoint.get("rng_state") if isinstance(rng_state, dict): restore_rng_state(rng_state, train_loader) else: resume_exact_rng = False discarded_records = len(records) - resume_epoch history_payload = previous_history history_payload["history"] = retained_records history_payload["completed_at"] = None history_payload["completed"] = False history_payload["epochs_requested"] = epochs history_payload["resume_exact_rng"] = resume_exact_rng resume_events = history_payload.setdefault("resume_events", []) resume_events.append( { "resumed_at": run_started, "checkpoint_path": str(resume_path), "checkpoint_epoch": resume_epoch, "discarded_history_records": discarded_records, "exact_rng_state_restored": resume_exact_rng, } ) best_metric = checkpoint_metric start_epoch = resume_epoch + 1 atomic_write_json(history_path, history_payload) print( f"resume checkpoint={resume_path} epoch={resume_epoch} " f"discarded_history_records={discarded_records} " f"exact_rng_state_restored={'yes' if resume_exact_rng else 'no'}", flush=True, ) print( f"device={device} parameters={parameter_count} " f"paper_parameters={paper_parameter_count} delta={parameter_count-paper_parameter_count}", flush=True, ) if device.type == "cuda" and torch.cuda.is_available(): torch.cuda.reset_peak_memory_stats(device) for epoch in range(start_epoch, epochs + 1): synchronize(device) epoch_start = time.perf_counter() lr_used = float(optimizer.param_groups[0]["lr"]) train_metrics = run_epoch(model, train_loader, device, epsilon, optimizer) test_metrics = run_epoch(model, test_loader, device, epsilon, optimizer=None) synchronize(device) duration = time.perf_counter() - epoch_start peak_memory_bytes = ( int(torch.cuda.max_memory_allocated(device)) if device.type == "cuda" and torch.cuda.is_available() else 0 ) monitor = float(train_metrics["full_relative_l2"]) is_best = monitor < best_metric if is_best: best_metric = monitor checkpoint = { "schema_version": "fno-ns2d-checkpoint-v1", "epoch": epoch, "model_state_dict": model.state_dict(), "optimizer_state_dict": optimizer.state_dict(), "scheduler_state_dict": scheduler.state_dict(), "best_train_full_relative_l2": best_metric, "test_full_relative_l2_at_best_epoch": float( test_metrics["full_relative_l2"] ), "monitor": "train_full_relative_l2", "test_selected": False, "config": deepcopy(config), "seed": seed, "parameter_count": parameter_count, "paper_parameter_count": paper_parameter_count, "parameter_count_difference": parameter_count - paper_parameter_count, "run_type": args.run_type, "scheduler_step_applied": False, "rng_state": capture_rng_state(train_loader), } atomic_torch_save(checkpoint_path, checkpoint) epoch_record = { "epoch": epoch, "learning_rate": lr_used, "duration_seconds": duration, "train_step_loss_sum": float(train_metrics["step_loss_sum"]), "train_mean_step_relative_l2": float( train_metrics["mean_step_relative_l2"] ), "train_full_relative_l2": float(train_metrics["full_relative_l2"]), "test_mean_step_relative_l2": float(test_metrics["mean_step_relative_l2"]), "test_full_relative_l2": float(test_metrics["full_relative_l2"]), "peak_accelerator_memory_bytes": peak_memory_bytes, "best": is_best, } history_payload["history"].append(epoch_record) history_payload["best_epoch"] = next( record["epoch"] for record in reversed(history_payload["history"]) if record["best"] ) history_payload["best_train_full_relative_l2"] = best_metric atomic_write_json(history_path, history_payload) print( f"epoch={epoch:04d}/{epochs:04d} time={duration:.3f}s lr={lr_used:.6g} " f"train_step_loss_sum={train_metrics['step_loss_sum']:.8f} " f"train_step_rel_l2={train_metrics['mean_step_relative_l2']:.8f} " f"train_full_rel_l2={train_metrics['full_relative_l2']:.8f} " f"test_step_rel_l2={test_metrics['mean_step_relative_l2']:.8f} " f"test_full_rel_l2={test_metrics['full_relative_l2']:.8f} " f"peak_mem_gib={peak_memory_bytes / (1024 ** 3):.3f} " f"best={'yes' if is_best else 'no'}", flush=True, ) scheduler.step() latest_checkpoint = { "schema_version": "fno-ns2d-training-state-v1", "epoch": epoch, "model_state_dict": model.state_dict(), "optimizer_state_dict": optimizer.state_dict(), "scheduler_state_dict": scheduler.state_dict(), "best_train_full_relative_l2": best_metric, "latest_train_full_relative_l2": float(train_metrics["full_relative_l2"]), "latest_test_full_relative_l2": float(test_metrics["full_relative_l2"]), "monitor": "train_full_relative_l2", "test_selected": False, "config": deepcopy(config), "seed": seed, "parameter_count": parameter_count, "paper_parameter_count": paper_parameter_count, "parameter_count_difference": parameter_count - paper_parameter_count, "run_type": args.run_type, "scheduler_step_applied": True, "rng_state": capture_rng_state(train_loader), } atomic_torch_save(latest_checkpoint_path, latest_checkpoint) history_payload["completed_at"] = datetime.now(timezone.utc).isoformat() history_payload["completed"] = True atomic_write_json(history_path, history_payload) print( f"training_complete best_epoch={history_payload['best_epoch']} " f"best_train_full_rel_l2={best_metric:.8f} checkpoint={checkpoint_path}", flush=True, ) if __name__ == "__main__": main()