| |
| """4000 full-objective Adam updates; every frame contributes to every update.""" |
| from __future__ import annotations |
|
|
| import argparse |
| import csv |
| import json |
| import os |
| import platform |
| 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 scripts.common import (ROOT, atomic_checkpoint, build_model, gather, load_dns, |
| predict, read_config, relative_l2, resolve, sample_indices, |
| save_json, setup, tensor) |
|
|
|
|
| def objective(model, boundary_q, boundary_uv, residual_q, cfg, backward): |
| """Exact sample means; chunking changes only summation order, not optimizer steps.""" |
| chunk = cfg["training"]["gradient_chunk_size"] |
| data_value, residual_value = 0.0, 0.0 |
| flow = model.flow_taylor if cfg["training"]["derivative_backend"] == "taylor" else model.flow |
| residual = model.residual_taylor if cfg["training"]["derivative_backend"] == "taylor" else model.residual_autograd |
| for start in range(0, len(boundary_q), chunk): |
| q = boundary_q[start:start + chunk] |
| uv = boundary_uv[start:start + chunk] |
| loss = ((flow(q)[:, :2] - uv)**2).sum() / len(boundary_q) |
| if backward: |
| loss.backward() |
| data_value += float(loss.detach()) |
| for start in range(0, len(residual_q), chunk): |
| _, r, _ = residual(residual_q[start:start + chunk]) |
| loss = (r**2).sum() / len(residual_q) |
| if backward: |
| loss.backward() |
| residual_value += float(loss.detach()) |
| regularization = cfg["training"]["slope_weight"] * model.slope_recovery() |
| if backward: |
| regularization.backward() |
| reg_value = float(regularization.detach()) |
| return {"loss": data_value + residual_value + reg_value, "loss_boundary": data_value, |
| "loss_residual": residual_value, "loss_slope": reg_value} |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--config", default=str(ROOT / "config/config.yaml")) |
| parser.add_argument("--device") |
| parser.add_argument("--resume", type=Path) |
| args = parser.parse_args() |
| cfg = read_config(args.config) |
| device, dtype = setup(cfg, args.device) |
| dns = load_dns(cfg) |
| out = resolve(cfg["outputs"]["results"]) |
| out.mkdir(parents=True, exist_ok=True) |
| indices = sample_indices(cfg, dns) |
| np.savez_compressed(out / "sampling_indices.npz", **indices) |
| bq, buv = gather(dns, indices["boundary"]) |
| rq, _ = gather(dns, indices["residual"]) |
| vq, vuv = gather(dns, indices["validation"]) |
| bq, buv, rq = (tensor(x, device, dtype) for x in (bq, buv, rq)) |
| model = build_model(cfg, device) |
| tr = cfg["training"] |
| optimizer = torch.optim.Adam(model.parameters(), lr=tr["learning_rate"], |
| betas=tuple(tr["adam_betas"]), eps=tr["adam_eps"]) |
| start_step, best_loss, best_step, elapsed_prior = 0, float("inf"), -1, 0.0 |
| if args.resume: |
| checkpoint = torch.load(args.resume, map_location=device, weights_only=False) |
| if checkpoint["config"] != cfg or checkpoint["data_sha256"] != dns["sha256"]: |
| raise ValueError("Resume requires identical config and DNS SHA256") |
| model.load_state_dict(checkpoint["model_state_dict"]) |
| optimizer.load_state_dict(checkpoint["optimizer_state_dict"]) |
| start_step = checkpoint["completed_updates"] |
| best_loss, best_step = checkpoint["best_loss"], checkpoint["best_step"] |
| elapsed_prior = checkpoint["elapsed_seconds"] |
| torch.set_rng_state(checkpoint["torch_rng_state"].cpu()) |
| if device.type == "cuda" and checkpoint.get("cuda_rng_state") is not None: |
| torch.cuda.set_rng_state(checkpoint["cuda_rng_state"].cpu()) |
| run_start = time.perf_counter() |
| parameter_count = sum(p.numel() for p in model.parameters()) |
| manifest = {"status": "running", "config": cfg, "data_path": dns["path"], "data_sha256": dns["sha256"], |
| "data_shapes": cfg["data"]["expected_shapes"], "time_range": [float(dns["t"][0]), float(dns["t"][-1])], |
| "time_spacing": float(np.diff(dns["t"])[0]), "boundary_samples": len(bq), "residual_samples": len(rq), |
| "sample_counts_per_time": {"residual": int(indices["residual"].shape[1]), "boundary": int(indices["boundary"].shape[1])}, |
| "validation_samples": len(vq), "parameter_count": parameter_count, |
| "framework": "PyTorch", "torch_version": torch.__version__, "python_version": platform.python_version(), |
| "device": str(device), "accelerator": torch.cuda.get_device_name(device) if device.type == "cuda" else "CPU", |
| "slurm_job_id": os.environ.get("SLURM_JOB_ID"), "dtype": str(dtype), |
| "derivative_backend": tr["derivative_backend"], "selection": tr["selection"], |
| "completed_updates": start_step, "started_unix": time.time()} |
| save_json(out / "run_manifest.json", manifest) |
| (out / "config_used.yaml").write_text(Path(args.config).read_text()) |
| print(json.dumps(manifest, ensure_ascii=False), flush=True) |
|
|
| def payload(step, score): |
| return {"format_version": 1, "class_name": "TSAPINN", "model_state_dict": model.state_dict(), |
| "optimizer_state_dict": optimizer.state_dict(), "config": cfg, |
| "completed_updates": step, "training_objective": score, |
| "best_loss": best_loss, "best_step": best_step, "data_sha256": dns["sha256"], |
| "elapsed_seconds": elapsed_prior + time.perf_counter() - run_start, |
| "torch_rng_state": torch.get_rng_state(), |
| "cuda_rng_state": torch.cuda.get_rng_state() if device.type == "cuda" else None} |
|
|
| history_path = out / "training_history.csv" |
| evaluation_path = out / "evaluation_history.csv" |
| history_fields = ["completed_updates", "loss", "loss_boundary", "loss_residual", "loss_slope", "learning_rate", "elapsed_seconds", "frequency_min", "frequency_mean", "frequency_max"] |
| mode = "a" if start_step else "w" |
| with history_path.open(mode, newline="", buffering=1) as hist, evaluation_path.open(mode, newline="", buffering=1) as evaluations: |
| writer = csv.DictWriter(hist, fieldnames=history_fields) |
| ew = csv.DictWriter(evaluations, fieldnames=["completed_updates", "heldout_u", "heldout_v", "snapshot_u", "snapshot_v"]) |
| if not start_step: |
| writer.writeheader() |
| ew.writeheader() |
| final_score = None |
| for completed in range(start_step, tr["iterations"] + 1): |
| |
| is_final = completed == tr["iterations"] |
| model.train() |
| optimizer.zero_grad(set_to_none=True) |
| with torch.set_grad_enabled(not is_final): |
| score = objective(model, bq, buv, rq, cfg, backward=not is_final) |
| if not all(np.isfinite(list(score.values()))): |
| raise FloatingPointError(f"Nonfinite objective at update {completed}: {score}") |
| if not is_final and any(p.grad is not None and not torch.isfinite(p.grad).all() for p in model.parameters()): |
| raise FloatingPointError(f"Nonfinite gradients at update {completed}") |
| if score["loss"] < best_loss: |
| best_loss, best_step = score["loss"], completed |
| atomic_checkpoint(resolve(cfg["outputs"]["checkpoint"]), payload(completed, score)) |
| elapsed = elapsed_prior + time.perf_counter() - run_start |
| frequencies = torch.cat([f.detach() for f in model.frequencies]) |
| lr_exponent = completed // tr["decay_steps"] if tr["staircase"] else completed / tr["decay_steps"] |
| lr = tr["learning_rate"] * tr["decay_rate"]**lr_exponent |
| for group in optimizer.param_groups: |
| group["lr"] = lr |
| row = {"completed_updates": completed, **score, "learning_rate": lr, "elapsed_seconds": elapsed, |
| "frequency_min": float(frequencies.min()), "frequency_mean": float(frequencies.mean()), "frequency_max": float(frequencies.max())} |
| writer.writerow(row) |
| if completed % tr["log_every"] == 0 or is_final: |
| print(f"update={completed:4d}/4000 loss={score['loss']:.8e} boundary={score['loss_boundary']:.8e} residual={score['loss_residual']:.8e} slope={score['loss_slope']:.8e} lr={lr:.8e} elapsed={elapsed:.2f}s", flush=True) |
| if completed % tr["eval_every"] == 0 or is_final: |
| val_prediction = predict(model, vq, device, dtype, cfg["evaluation"]["batch_size"]) |
| val_error = relative_l2(val_prediction[:, :2], vuv) |
| snap = int(np.argmin(abs(dns["t"] - cfg["evaluation"]["snapshot_time"]))) |
| sq = np.column_stack((dns["xy"], np.full(len(dns["xy"]), dns["t"][snap]))) |
| spred = predict(model, sq, device, dtype, cfg["evaluation"]["batch_size"]) |
| se = relative_l2(spred[:, :2], dns["uv"][snap]) |
| ew.writerow({"completed_updates": completed, "heldout_u": val_error[0], "heldout_v": val_error[1], "snapshot_u": se[0], "snapshot_v": se[1]}) |
| print(f"EVAL update={completed} heldout_relL2(u,v)=({val_error[0]:.8e},{val_error[1]:.8e}) snapshot_t={dns['t'][snap]:.3f} relL2(u,v)=({se[0]:.8e},{se[1]:.8e})", flush=True) |
| manifest.update(completed_updates=completed, best_step=best_step, best_loss=best_loss, elapsed_seconds=elapsed) |
| save_json(out / "run_manifest.json", manifest) |
| if completed % tr["checkpoint_every"] == 0 or is_final: |
| atomic_checkpoint(resolve(cfg["outputs"]["last_checkpoint"]), payload(completed, score)) |
| if is_final: |
| final_score = score |
| break |
| optimizer.step() |
| elapsed = elapsed_prior + time.perf_counter() - run_start |
| manifest.update(status="completed", completed_updates=tr["iterations"], elapsed_seconds=elapsed, |
| best_step=best_step, best_loss=best_loss, final_objective=final_score, |
| peak_device_memory_bytes=torch.cuda.max_memory_allocated() if device.type == "cuda" else 0) |
| save_json(out / "run_manifest.json", manifest) |
| print(f"TRAINING COMPLETE updates=4000 best_step={best_step} best_loss={best_loss:.8e} runtime_seconds={elapsed:.3f}", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
