diff options
Diffstat (limited to 'src/nn_train.py')
| -rwxr-xr-x | src/nn_train.py | 58 |
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) |