flood fill for roads, fix dragon, memoise leaf eval, enable diagonal in nn_train.py

1 files changed, 27 insertions, 31 deletions

diff --git a/src/nn_train.py b/src/nn_train.py
index 2eb9b0c..8e648e9 100755
--- a/src/nn_train.py
+++ b/src/nn_train.py
@@ -357,6 +357,7 @@ class Net(Module):
     ) -> None:
         upper = "FIVE" if size == 5 else "SIX"
         prefix = (f"cand_{id}_" if id is not None else "") + upper
+        guard = f"WEIGHTS_{upper}" + (f"_{id}" if id is not None else "")
 
         E = self.embedding.weight.detach().cpu().numpy()
         W = self.hidden.weight.detach().cpu().numpy()
@@ -365,26 +366,20 @@ class Net(Module):
         LB = float(self.output.bias.detach().cpu().item())
 
         with open(path, "w") as f:
-            _ = f.write("/* generated by nn_train.py, do not edit by hand */\n")
-            _ = f.write(f"/* {note} */\n")
-            _ = f.write(f'#include "weights_{size}.h"\n\n')
-            _ = f.write(
-                f"const float {prefix}_embedding[{len(E)}] =\n"
-                f"{{\n{_fmt_flat(E, per_line=7)}\n}};\n\n"
-            )
-            _ = f.write(
-                f"const float {prefix}_dense_weights[{upper}_DENSE_NUM]"
-                f"[{upper}_INP_NUM] NN_ALIGNED =\n{{\n{_fmt_matrix(W)}\n}};\n\n"
-            )
-            _ = f.write(
-                f"const float {prefix}_dense_biases[{upper}_DENSE_NUM] =\n"
-                f"{{\n{_fmt_flat(B)}\n}};\n\n"
-            )
-            _ = f.write(
-                f"const float {prefix}_logit_weights[{upper}_DENSE_NUM]"
-                f" NN_ALIGNED =\n{{\n{_fmt_flat(LW)}\n}};\n\n"
-            )
-            _ = f.write(f"const float {prefix}_logit_bias = {LB:.9g};\n")
+            lines = [
+                "/* generated by nn_train.py, do not edit by hand */\n",
+                f"/* {note} */\n",
+                f"#ifndef {guard}\n#define {guard}\n",
+                f'#include "weights_{size}.h"\n\n',
+                f"const float {prefix}_embedding[{len(E)}] =\n{{\n{_fmt_flat(E, per_line=7)}\n}};\n\n",
+                f"const float {prefix}_dense_weights[{upper}_DENSE_NUM][{upper}_INP_NUM] NN_ALIGNED =\n{{\n{_fmt_matrix(W)}\n}};\n\n",
+                f"const float {prefix}_dense_biases[{upper}_DENSE_NUM] =\n{{\n{_fmt_flat(B)}\n}};\n\n",
+                f"const float {prefix}_logit_weights[{upper}_DENSE_NUM] NN_ALIGNED =\n{{\n{_fmt_flat(LW)}\n}};\n\n",
+                f"const float {prefix}_logit_bias = {LB:.9g};\n",
+                f"#endif",
+            ]
+            for line in lines:
+                _ = f.write(line)
 
 
 # ===================================================================
@@ -431,6 +426,7 @@ def cand_copy_for(size: int, id: int, root: Path) -> None:
         "logit_bias",
     ):
         txt = txt.replace(f"{upper}_{sym}", f"cand_{id}_{upper}_{sym}")
+    txt = txt.replace(f"WEIGHTS_{upper}", f"WEIGHTS_{upper}_{id}")
     _ = (root / "include" / f"cand_{id}_weights_{size}.c").write_text(txt)
 
 
@@ -477,7 +473,7 @@ def run_step(
 
     num_candidates = len(candidates)
     cand_pairs = [
-        (i, j) for i in range(num_candidates) for j in range(i + 1, num_candidates)
+        (i, j) for i in range(num_candidates) for j in range(i, num_candidates)
     ]
 
     # 1. rebuild
@@ -488,7 +484,7 @@ def run_step(
 
     # 2. generate self-play data
     csv = csv_path(iter)
-    print(f"selfplay: {num_candidates=}, {size=}, {depth=}, {min_ply=}")
+    print(f"selfplay: {num_candidates=}, {games=}, {size=}, {depth=}, {min_ply=}")
     out_rows: list[str] = []
     for out in fan_out(
         [
@@ -519,7 +515,7 @@ def run_step(
     for idx, _, perf in sorted(results, key=lambda t: t[2].rate(), reverse=True):
         print(f"\t cand {idx:03d} {perf} {perf.rate() * 100.0:.2f}%")
 
-    # 4. fine-tune the champion on TD(lambda) targets -> candidate
+    # 4. fine-tune on TD(lambda) targets -> candidate
     for idx, cand in enumerate(candidates):
         corr, mse = cand.fit(
             records=records.index(idx),
@@ -537,17 +533,17 @@ def run_step(
         )
 
     # 5. write champ
-    _, champ, perf = max(results, key=lambda t: t[2].rate())
+    idx, champ, perf = max(results, key=lambda t: t[2].rate())
+    path = root / "include" / f"weights_{size}.c"
     champ.to_c(
         size=size,
         id=None,
-        path=root / "include" / f"weights_{size}.c",
-        note=f"{iter=}, perf={str(perf)}",
+        path=path,
+        note=f"{iter=}, perf={str(perf)}, rate={perf.rate() * 100.0:.3f}",
     )
     _ = subprocess.run(["make", "clean"], cwd=root, check=True, capture_output=True)
-    for i, j in cand_pairs:
-        (root / f"selfplay-{i}-{j}").unlink(missing_ok=True)
     _ = subprocess.run(["make", "native"], cwd=root, check=True, capture_output=True)
+    print(f"Wrote candidate {idx=} to {str(path)}")
 
 
 # ===================================================================
@@ -577,7 +573,7 @@ def main() -> None:
     _ = p.add_argument(
         "--games",
         type=int,
-        default=100,
+        default=200,
         help="selfplay games per iteration",
     )
     _ = p.add_argument(
@@ -626,7 +622,7 @@ def main() -> None:
         "--procs",
         type=int,
         default=None,
-        help="parallel game processes (half the CPUs if unset)",
+        help="parallel game processes (number of CPUs if unset)",
     )
     _ = p.add_argument(
         "--depth",
@@ -644,7 +640,7 @@ def main() -> None:
     args = p.parse_args()
 
     if args.procs is None:
-        args.procs = max(1, (os.cpu_count() or 1) // 2)
+        args.procs = max(1, (os.cpu_count() or 1))
 
     (root / "data").mkdir(exist_ok=True)
     candidates = setup(num_candidates=args.num_candidates, size=args.size, root=root)