File size: 4,597 Bytes
0f0d1f4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 | #!/usr/bin/env python3
"""Six-path smoke checks and independent derivative/gradient equivalence checks."""
from __future__ import annotations
import argparse
import json
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, build_model, gather, load_dns, predict, read_config,
relative_l2, sample_indices, save_json, setup, tensor)
from scripts.train import objective
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--device", default="cuda")
parser.add_argument("--benchmark", action="store_true")
args = parser.parse_args()
cfg = read_config()
device, dtype = setup(cfg, args.device)
dns = load_dns(cfg)
index = sample_indices(cfg, dns)
bq, buv = gather(dns, index["boundary"])
rq, _ = gather(dns, index["residual"])
model = build_model(cfg, device)
q = tensor(rq[:32], device, dtype)
checks = {}
output = model(q)
assert output.shape == (32, 2) and torch.isfinite(output).all()
checks["forward_pass"] = "passed"
values, residuals, div = model.residual_autograd(q)
tvalues, tresiduals, _ = model.residual_taylor(q)
torch.testing.assert_close(tvalues, values, atol=1e-10, rtol=1e-9)
torch.testing.assert_close(tresiduals, residuals, atol=1e-10, rtol=1e-9)
params = tuple(model.parameters())
aloss = residuals.square().mean() + values[:, :2].square().mean() + 0.25 * model.slope_recovery()
tloss = tresiduals.square().mean() + tvalues[:, :2].square().mean() + 0.25 * model.slope_recovery()
ag = torch.autograd.grad(aloss, params)
tg = torch.autograd.grad(tloss, params)
for left, right in zip(ag, tg):
torch.testing.assert_close(left, right, atol=2e-10, rtol=2e-8)
assert float(div.abs().max()) < 1e-10
checks["backward_pass"] = "passed: finite parameter gradients; Taylor/reverse-mode equivalence"
bq_t, buv_t, rq_t = (tensor(x, device, dtype) for x in (bq[:64], buv[:64], rq[:64]))
opt = torch.optim.Adam(model.parameters(), lr=cfg["training"]["learning_rate"], eps=1e-7)
losses = []
for _ in range(3):
opt.zero_grad(set_to_none=True)
value = objective(model, bq_t, buv_t, rq_t, cfg, backward=True)
opt.step()
losses.append(value["loss"])
assert np.isfinite(losses).all() and losses[-1] < losses[0], losses
checks["train_loop_dry_run"] = {"status": "passed", "losses": losses}
valq, valuv = gather(dns, index["validation"])
valp = predict(model, valq[:128], device, dtype, 64)
error = relative_l2(valp[:, :2], valuv[:128])
assert np.isfinite(error).all()
checks["validation_loop_dry_run"] = {"status": "passed", "relative_l2": error.tolist()}
assert np.allclose(relative_l2(dns["uv"].reshape(-1, 2), dns["uv"].reshape(-1, 2)), 0)
assert np.allclose(relative_l2(np.zeros_like(dns["uv"].reshape(-1, 2)), dns["uv"].reshape(-1, 2)), 1)
checks["domain_metric_dry_run"] = "passed: L2 self=0, zero predictor=1; streamfunction divergence <1e-10"
assert len(model.frequencies) == 4 and sum(f.numel() for f in model.frequencies) == 200
assert [l.out_features for l in model.layers] == [50, 50, 50, 50, 2]
assert index["boundary"].shape == (200, 400) and index["residual"].shape == (200, 1000)
checks["config_consistency"] = "passed: 4x50, 200 frequencies, (x,y,t)->(psi,p), 4000 configured updates, per-time counts"
benchmark = None
if args.benchmark:
model = build_model(cfg, device)
fullb, fulluv, fullr = (tensor(x, device, dtype) for x in (bq, buv, rq))
started = time.perf_counter()
model.zero_grad(set_to_none=True)
value = objective(model, fullb, fulluv, fullr, cfg, backward=True)
if device.type == "cuda":
torch.cuda.synchronize()
benchmark = {"seconds_one_full_objective_backward": time.perf_counter()-started, "score": value,
"peak_device_memory_bytes": torch.cuda.max_memory_allocated() if device.type == "cuda" else 0}
result = {"status": "passed", "attempt": 1, "checks": checks, "benchmark": benchmark,
"parameter_count": sum(p.numel() for p in model.parameters()), "torch": torch.__version__,
"device": torch.cuda.get_device_name() if device.type == "cuda" else "CPU"}
save_json(ROOT / ".onescience/smoke_result.json", result)
print(json.dumps(result, ensure_ascii=False, indent=2), flush=True)
if __name__ == "__main__":
main()
|