factor out nn_train.py into a separate python project to make it more manageable, add plotting and offline strength fit

1 files changed, 299 insertions, 0 deletions

diff --git a/train/src/ctak_train/nn_train.py b/train/src/ctak_train/nn_train.py
new file mode 100644
index 0000000..d2c2990
--- /dev/null
+++ b/train/src/ctak_train/nn_train.py
@@ -0,0 +1,299 @@
+#
+# This file is part of ctak.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with ctak. If not, see <https://www.gnu.org/licenses/>.
+
+import argparse
+import os
+import subprocess
+from pathlib import Path
+
+import numpy as np
+
+from ctak_train.candidate import Candidate, DataGenConfig, TrainingGridSearchConfig
+from ctak_train.data import GameLog, load_window
+from ctak_train.net import NET_DIMS
+from ctak_train.selfplay import fan_out
+from ctak_train.strength import fit_strengths, plot_strengths
+
+
+def run_step(
+    *,
+    candidates: list[Candidate],
+    data_gen_config: DataGenConfig,
+    log: GameLog,
+    strength_history: list[np.ndarray],
+    iter: int,
+):
+    def data_path(i: int) -> Path:
+        return (
+            data_gen_config.root
+            / "data"
+            / f"selfplay-{data_gen_config.size}-v{i:03d}.json"
+        )
+
+    print(f"\n== iteration {iter} ==")
+
+    num_candidates = len(candidates)
+    cand_pairs = [
+        (i, j) for i in range(num_candidates) for j in range(i, num_candidates)
+    ]
+
+    # 1. rebuild
+    for i, j in cand_pairs:
+        _ = subprocess.run(
+            ["make", f"selfplay-{i}-{j}"],
+            cwd=data_gen_config.root,
+            check=True,
+            capture_output=True,
+        )
+
+    # 2. generate self-play data
+    data = data_path(iter)
+    print(
+        f"selfplay: {num_candidates=}, {data_gen_config.games=}, "
+        f"{data_gen_config.size=}, {data_gen_config.depth=}, {data_gen_config.min_ply=}, "
+        f"{data_gen_config.train_only_self_play=}"
+    )
+    out_games: list[str] = []
+    for out in fan_out(
+        [
+            [
+                f"./selfplay-{i}-{j}",
+                str(data_gen_config.size),
+                str(data_gen_config.depth),
+                str(data_gen_config.min_ply),
+            ]
+            for (i, j) in cand_pairs
+        ],
+        data_gen_config.games,
+        data_gen_config.procs,
+        "selfplay",
+        data_gen_config.root,
+    ):
+        out_games.extend(out.splitlines())
+    print(f"writing {len(out_games):,} games to {data}")
+    with open(data, "w") as f:
+        _ = f.write("\n".join(out_games))
+        if out_games:
+            _ = f.write("\n")
+
+    # 3. load the training window and append this iteration's games to the log
+    data_files = [
+        data_path(i)
+        for i in range(max(iter - data_gen_config.replay_iters, 1), iter + 1)
+    ]
+    print("Loading data from", [str(c) for c in data_files], "...", end=" ")
+    pools, new_games = load_window(
+        data_files, data_gen_config.size, num_candidates, iter
+    )
+    print("candidates have", [f"{p.n_rows:,}" for p in pools], "rows.")
+    log.append(new_games)
+
+    # 4. fit strengths over the whole history, and plot the trajectory
+    strengths, se = fit_strengths(log, num_candidates)
+    strength_history.append(strengths)
+    plot_strengths(
+        strength_history,
+        candidates,
+        data_gen_config.root / "data" / "candidate_strength.png",
+    )
+
+    rated_games = int(np.count_nonzero(log.black != log.white))
+    ranked = sorted(range(num_candidates), key=lambda i: strengths[i], reverse=True)
+    print(
+        f"Strength fit ({rated_games:,} head-to-head games, "
+        f"iterations 1..{iter}, {rated_games // num_candidates:,} games/candidate):"
+    )
+    for idx in ranked:
+        print(
+            f"\tcandidate {idx:03d}: strength {strengths[idx]:6.1f} "
+            f"± {1.96 * se[idx]:5.1f} Config {candidates[idx]}"
+        )
+
+    # 5. write champ, unconditionally
+    idx = ranked[0]
+    path = data_gen_config.root / "include" / f"weights_{data_gen_config.size}.c"
+    candidates[idx].net.to_c(
+        size=data_gen_config.size,
+        ident=None,
+        path=path,
+        note=f"{iter=}, strength={strengths[idx]:.1f}",
+    )
+    _ = subprocess.run(
+        ["make", "clean"], cwd=data_gen_config.root, check=True, capture_output=True
+    )
+    _ = subprocess.run(
+        ["make", "native"], cwd=data_gen_config.root, check=True, capture_output=True
+    )
+    print(f"Wrote candidate {idx:03d} to {path!s}")
+
+    # 6. fine-tune on the hybrid targets -> candidate
+    for idx, cand in enumerate(candidates):
+        corr, mse = cand.fit(pools[idx])
+        cand.net.to_c(
+            size=data_gen_config.size,
+            ident=idx,
+            path=data_gen_config.root
+            / "include"
+            / f"cand_{idx}_weights_{data_gen_config.size}.c",
+            note=f"id={idx}, v{iter:03d}, val_corr={corr:.4f}, val_mse={mse:.4f}",
+        )
+
+
+def main() -> None:
+    # repo root, relative to this file at <root>/train/src/ctak_train/nn_train.py
+    root = Path(__file__).resolve().parents[3]
+    p = argparse.ArgumentParser(
+        description="Train the neural network evaluators of ctak",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    _ = p.add_argument(
+        "--size",
+        type=int,
+        choices=sorted(NET_DIMS),
+        default=5,
+        help="board size",
+    )
+    _ = p.add_argument(
+        "--iters",
+        type=int,
+        default=100,
+        help="max loop iterations",
+    )
+    _ = p.add_argument(
+        "--games",
+        type=int,
+        default=200,
+        help="games per iteration",
+    )
+    _ = p.add_argument(
+        "--replay-iters",
+        type=int,
+        default=10,
+        help="number of previous data files to keep in the training pools "
+        "(every game record is kept for the strength fit)",
+    )
+    _ = p.add_argument(
+        "--smooth-factor",
+        type=float,
+        nargs="+",
+        default=[0.3, 0.5, 0.7, 0.9],
+        help="grid points for the EMA smoothing factor",
+    )
+    _ = p.add_argument(
+        "--epochs",
+        type=int,
+        nargs="+",
+        default=[5],
+        help="grid points for candidate training epochs",
+    )
+    _ = p.add_argument(
+        "--lr",
+        type=float,
+        nargs="+",
+        default=[1e-4],
+        help="grid points for learning rate",
+    )
+    _ = p.add_argument(
+        "--batch",
+        type=int,
+        nargs="+",
+        default=[128],
+        help="grid points for training batch size",
+    )
+    _ = p.add_argument(
+        "--val-frac",
+        type=float,
+        nargs="+",
+        default=[0.05],
+        help="grid points for fraction of rows held out for validation",
+    )
+    _ = p.add_argument(
+        "--procs",
+        type=int,
+        default=None,
+        help="parallel game processes (number of CPUs if unset)",
+    )
+    _ = p.add_argument(
+        "--depth",
+        type=int,
+        default=4,
+        help="search depth for selfplay and the matches",
+    )
+    _ = p.add_argument(
+        "--min-ply",
+        type=int,
+        default=8,
+        help="length of random opening in games",
+    )
+    _ = p.add_argument(
+        "--train-all-games",
+        action="store_true",
+        help="train on all selfplay games, including matches between "
+        "different candidates (default: self vs self games only)",
+    )
+
+    args = p.parse_args()
+
+    # avoid counter-play start
+    if args.min_ply < 2:
+        _ = p.error("--min-ply must be >= 2")
+
+    if args.procs is None:
+        args.procs = max(1, (os.cpu_count() or 1))
+
+    data_gen = DataGenConfig(
+        size=args.size,
+        depth=args.depth,
+        min_ply=args.min_ply,
+        games=args.games,
+        procs=args.procs,
+        replay_iters=args.replay_iters,
+        root=root,
+        train_only_self_play=not args.train_all_games,
+    )
+    tgsc = TrainingGridSearchConfig(
+        smooth_factor=args.smooth_factor,
+        epochs=args.epochs,
+        batch=args.batch,
+        lr=args.lr,
+        val_frac=args.val_frac,
+    )
+
+    candidates = tgsc.make_configs(data_gen)
+
+    (root / "data").mkdir(exist_ok=True)
+
+    num_params = sum(
+        p.numel() for p in candidates[0].net.parameters() if p.requires_grad
+    )
+    print(f"Number of trainable parameters: {num_params:,}")
+
+    log = GameLog.empty()
+    strength_history: list[np.ndarray] = []
+
+    for iter in range(1, args.iters + 1):
+        run_step(
+            candidates=candidates,
+            data_gen_config=data_gen,
+            log=log,
+            strength_history=strength_history,
+            iter=iter,
+        )
+    print(f"\n== done, champion in include/weights_{args.size}.c ==")
+
+
+if __name__ == "__main__":
+    main()