"""Formal Table 3 training: 20,000 Adam updates, batch64 per constraint.""" import argparse import csv import fcntl import json import random import resource import sys import time from pathlib import Path sys.path.insert(0, str(Path(__file__).resolve().parents[1])) import numpy as np import torch from common import (ROOT, atomic_json, atomic_save, build_model, evaluate, load_checkpoint, load_config, load_validation, loss_terms, make_sobol_pools, require_compute_node, resolved, sample_batches, seed_everything, selection_score, sha256) def train(config, seed, resume=None): allocation = require_compute_node() if seed not in config["training"]["seeds"]: raise ValueError("Seed is not in the formal experiment list") tcfg = config["training"] seed_everything(seed, tcfg["num_threads"]) out = resolved(config, "results") / f"seed_{seed}" weights = resolved(config, "weights") / f"seed_{seed}" out.mkdir(parents=True, exist_ok=True) weights.mkdir(parents=True, exist_ok=True) if (out/"history.csv").exists() and resume is None: raise FileExistsError("Existing training history; use --resume to preserve run provenance") model = build_model(config) pools = make_sobol_pools(config, seed) rng = np.random.default_rng(seed+100000) xval, yval, valids = load_validation(config) np.save(out/"validation_indices.npy", valids) optimizer = torch.optim.Adam(model.parameters(), lr=tcfg["learning_rate"], betas=tuple(tcfg["betas"]), eps=tcfg["epsilon"], weight_decay=tcfg["weight_decay"]) sched = tcfg["scheduler"] scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode="min", factor=sched["factor"], patience=sched["patience"], threshold=sched["threshold"]) start_epoch, best_score, best_epoch = 1, float("inf"), None accumulated_seconds = 0.0 if resume: previous = load_checkpoint(resume, model) if previous["config"] != config or previous["seed"] != seed: raise ValueError("Resume config/seed differs from checkpoint") optimizer.load_state_dict(previous["optimizer_state_dict"]) scheduler.load_state_dict(previous["scheduler_state_dict"]) start_epoch = previous["epoch"]+1 best_score, best_epoch = previous["best_score"], previous["best_epoch"] torch.set_rng_state(previous["torch_rng"]) random.setstate(previous["python_rng"]) rng.bit_generator.state = previous["batch_rng"] accumulated_seconds = previous["elapsed_seconds"] provenance = {"config": config, "seed": seed, "allocation": allocation, "reference": str(resolved(config, "reference")), "reference_sha256": sha256(resolved(config, "reference")), "parameter_count": sum(p.numel() for p in model.parameters()), "torch_version": torch.__version__, "epoch_semantics": "one Adam step, 64 points per constraint", "reference_used_for_gradient_training": False, "validation_indices": str(out/"validation_indices.npy")} atomic_json(out/"run_manifest.json", provenance) print(f"START seed={seed} allocation={json.dumps(allocation)} parameters=923 epochs=20000 batch=64", flush=True) columns = ["epoch", "total", "physics", "lid", "walls", "initial", "lr", "gradient_norm", "step_seconds", "elapsed_seconds", "val_u_l2", "val_v_l2", "val_p_l2", "val_mean_l2"] mode = "a" if resume else "w" start = time.monotonic() with open(out/"history.csv", mode, newline="", buffering=1) as history: writer = csv.DictWriter(history, fieldnames=columns) if not resume: writer.writeheader() for epoch in range(start_epoch, tcfg["epochs"]+1): tick = time.monotonic() model.train() optimizer.zero_grad(set_to_none=True) terms = loss_terms(model, sample_batches(pools, tcfg["batch_size"], rng), config) total = terms["total"] if not torch.isfinite(total): raise FloatingPointError(f"Nonfinite loss at {epoch}") total.backward() grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), tcfg["gradient_clip"], error_if_nonfinite=True) optimizer.step() values = {k: float(v.detach()) for k, v in terms.items()} scheduler.step(values["total"]) step_seconds = time.monotonic()-tick elapsed = accumulated_seconds+time.monotonic()-start row = {"epoch": epoch, **values, "lr": optimizer.param_groups[0]["lr"], "gradient_norm": float(grad_norm), "step_seconds": step_seconds, "elapsed_seconds": elapsed} metrics = None if epoch == 1 or epoch % config["validation"]["every"] == 0 or epoch == tcfg["epochs"]: metrics = evaluate(model, xval, yval, config["validation"]["batch_size"]) score = selection_score(metrics) row.update({f"val_{k}_l2": metrics[k]["relative_l2"] for k in ("u", "v", "p")}) row["val_mean_l2"] = score if score < best_score: best_score, best_epoch = score, epoch is_best = True else: is_best = False state = {"model_state_dict": model.state_dict(), "config": config, "seed": seed, "epoch": epoch, "validation_metrics": metrics, "selection_score": score, "best_score": best_score, "best_epoch": best_epoch, "optimizer_state_dict": optimizer.state_dict(), "scheduler_state_dict": scheduler.state_dict(), "torch_rng": torch.get_rng_state(), "python_rng": random.getstate(), "batch_rng": rng.bit_generator.state, "elapsed_seconds": elapsed, "allocation": allocation, "loss": values} if is_best: atomic_save(weights/"best_model.pt", state) # Publish the validation-best checkpoint at the requested root path # throughout concurrent training, with one writer at a time. root_weights = resolved(config, "weights") root_meta = resolved(config, "results")/"best_checkpoint_selection.json" with open(root_weights/".best_model.lock", "a") as lock: try: fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB) except BlockingIOError: # Another seed is publishing; per-seed checkpoint remains authoritative. pass else: incumbent = json.loads(root_meta.read_text()) if root_meta.exists() else None if incumbent is None or score < incumbent["selection_score"]: atomic_save(root_weights/"best_model.pt", state) atomic_json(root_meta, {"seed": seed, "epoch": epoch, "selection_score": score, "metrics": metrics, "checkpoint": str(root_weights/"best_model.pt")}) fcntl.flock(lock, fcntl.LOCK_UN) if epoch % tcfg["checkpoint_every"] == 0 or epoch == tcfg["epochs"]: atomic_save(weights/"last_model.pt", state) atomic_json(out/"progress.json", {"status": "training", "epoch": epoch, "epochs": tcfg["epochs"], "loss": values, "validation_metrics": metrics, "best_score": best_score, "best_epoch": best_epoch, "allocation": allocation}) writer.writerow(row) line = f"seed={seed} epoch={epoch}/20000 loss={values['total']:.8e} phy={values['physics']:.6e} lid={values['lid']:.6e} walls={values['walls']:.6e} ic={values['initial']:.6e} lr={row['lr']:.6e}" if metrics is not None: line += " " + " ".join(f"{k}_relL2={metrics[k]['relative_l2']:.8e} ({metrics[k]['relative_l2_percent']:.5f}%)" for k in ("u", "v", "p")) print(line, flush=True) elapsed = accumulated_seconds + time.monotonic()-start atomic_json(out/"training_summary.json", {"status": "completed", "completed_epochs": tcfg["epochs"], "seed": seed, "best_epoch": best_epoch, "best_validation_mean_l2": best_score, "last_loss": values, "elapsed_seconds": elapsed, "peak_rss_mb": resource.getrusage(resource.RUSAGE_SELF).ru_maxrss/1024, "best_checkpoint": str(weights/"best_model.pt"), "last_checkpoint": str(weights/"last_model.pt"), "allocation": allocation}) atomic_json(out/"progress.json", {"status": "completed", "epoch": tcfg["epochs"], "best_epoch": best_epoch}) print(f"TRAINING_COMPLETED seed={seed} epochs=20000 best_epoch={best_epoch}", flush=True) if __name__ == "__main__": parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--config", default=str(ROOT/"config/config.yaml")) parser.add_argument("--seed", type=int, required=True) parser.add_argument("--resume", type=Path) args = parser.parse_args() train(load_config(args.config), args.seed, args.resume)