| """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) |
| |
| |
| 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: |
| |
| 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) |
|
|
