aboutsummaryrefslogtreecommitdiff
path: root/src/nn_train.py
diff options
context:
space:
mode:
Diffstat (limited to 'src/nn_train.py')
-rwxr-xr-x[-rw-r--r--]src/nn_train.py643
1 files changed, 564 insertions, 79 deletions
diff --git a/src/nn_train.py b/src/nn_train.py
index 03852f5..52549d8 100644..100755
--- a/src/nn_train.py
+++ b/src/nn_train.py
@@ -1,6 +1,24 @@
-# cnn_train.py, train a small CNN to recognise winning Tak positions
+#!/usr/bin/env -S uv run --script
#
-# Copyright (C) 2021, tslil clingman
+# /// script
+# requires-python = ">=3.14"
+# dependencies = [
+# "pandas>=3.0.5",
+# "torch>=2.14.0",
+# "tqdm>=4.70.0",
+# ]
+#
+# [tool.uv.sources]
+# torch = { index = "pytorch-cpu" }
+#
+# [[tool.uv.index]]
+# name = "pytorch-cpu"
+# url = "https://download.pytorch.org/whl/cpu"
+# explicit = true
+# ///
+
+#
+# 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
@@ -13,95 +31,562 @@
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
-# along with this program. If not, see <https://www.gnu.org/licenses/>.
+# along with ctak. If not, see <https://www.gnu.org/licenses/>.
+from collections.abc import Iterator
+from dataclasses import dataclass
+
+import argparse
+import concurrent.futures
+import copy
+import os
+from pathlib import Path
+import subprocess
-# Loading data
-import pandas as pd
import numpy as np
+import pandas as pd
+import torch
+from tqdm import tqdm
+
+from torch.nn import MSELoss, Linear, Module
+from torch import Tensor
+
+
+# ===================================================================
+# Net dimensions
+# ===================================================================
+
+# (input dim, hidden dim) per board size; must match include/weights_<size>.h
+NET_DIMS: dict[int, tuple[int, int]] = {
+ 5: (28, 64), # FIVE_INP_NUM, FIVE_DENSE_NUM
+ 6: (39, 64), # SIX_INP_NUM, SIX_DENSE_NUM
+}
+
+
+# ===================================================================
+# C weight formatting
+# ===================================================================
+
+
+def _fmt_flat(a: np.ndarray, per_line: int = 8) -> str:
+ a = np.ravel(np.asarray(a, dtype=np.float32))
+ lines: list[str] = []
+ for i in range(0, len(a), per_line):
+ chunk = ", ".join(f"{v:.9g}" for v in a[i : i + per_line])
+ if i + per_line < len(a):
+ chunk += ","
+ lines.append(" " + chunk)
+ return "\n".join(lines)
+
+
+def _fmt_matrix(mat: np.ndarray, per_line: int = 8) -> str:
+ mat = np.asarray(mat, dtype=np.float32)
+ rows: list[str] = []
+ for r in mat:
+ chunks = [
+ ", ".join(f"{v:.9g}" for v in r[i : i + per_line])
+ for i in range(0, len(r), per_line)
+ ]
+ rows.append(" {" + ",\n ".join(chunks) + "},")
+ return "\n".join(rows)
+
+
+# ===================================================================
+# Self-play records
+# ===================================================================
+
+
+def td_lambda_targets(
+ values: Tensor,
+ result: float,
+ td_lambda: float,
+) -> Tensor:
+ # values: V(s_t) for t = 0..T-1, shape [T]
+ # result: final result R of the game
+
+ # This is ema, computed backwards
+ # target[T-1] = R
+ # target[t] = (1 - lambda) * V(s_{t+1}) + lambda * target[t+1]
+ length = len(values)
+ targets = torch.empty_like(values)
+ targets[-1] = result
+ for ply in range(length - 2, -1, -1):
+ targets[ply] = (1.0 - td_lambda) * values[ply + 1] + td_lambda * targets[
+ ply + 1
+ ]
+ return targets
+
+
+@dataclass(frozen=True)
+class SelfPlayRecords:
+ positions: Tensor # [n, inp_dim] feature vector of each position
+ game_ids: Tensor # [n] id of the game the position belongs to
+ plies: Tensor # [n] ply number (position is before this move)
+ results: Tensor # [n] final result from black's perspective (+1 black won, -1 white won, 0 draw)
+
+ @classmethod
+ def load_selfplay(
+ cls, fn: Path, size: int, device: torch.device
+ ) -> SelfPlayRecords:
+ # [ net input: 0-inp_dim ] [game_id] [ply] [game result]
+ df = pd.read_csv(fn, header=None)
+ inp_dim = NET_DIMS[size][0]
+ return cls(
+ positions=torch.tensor(
+ df.iloc[:, 0:inp_dim].to_numpy(dtype=np.float32),
+ dtype=torch.float32,
+ device=device,
+ ),
+ game_ids=torch.tensor(
+ df.iloc[:, inp_dim].to_numpy(dtype=np.int64),
+ device=device,
+ ),
+ plies=torch.tensor(
+ df.iloc[:, inp_dim + 1].to_numpy(dtype=np.int64),
+ device=device,
+ ),
+ results=torch.tensor(
+ df.iloc[:, inp_dim + 2].to_numpy(dtype=np.float32),
+ dtype=torch.float32,
+ device=device,
+ ),
+ )
+
+ def make_training_data(
+ self, model: Module, td_lambda: float
+ ) -> tuple[Tensor, Tensor]:
+ model.eval()
+ with torch.no_grad():
+ values = model(self.positions)
+
+ targets = torch.empty_like(values)
+
+ for game_id in torch.unique(self.game_ids):
+ in_game = self.game_ids == game_id
+ game_values = values[in_game]
+ game_result = self.results[in_game][0].item()
+ targets[in_game] = td_lambda_targets(game_values, game_result, td_lambda)
+
+ return self.positions, targets
+
+
+# ===================================================================
+# Network
+# ===================================================================
+
+
+class Net(Module):
+ def __init__(self, inp_dim: int, hidden: int) -> None:
+ super().__init__()
+ self.hidden = Linear(inp_dim, hidden, bias=True)
+ self.output = Linear(hidden, 1, bias=True)
+
+ def forward(self, x: Tensor) -> Tensor:
+ hidden = torch.relu(self.hidden(x))
+ return torch.tanh(self.output(hidden)).squeeze(-1)
+
+ def fit(
+ self,
+ *,
+ self_play_records: SelfPlayRecords,
+ td_lambda: float,
+ val_frac: float,
+ epochs: int,
+ batch: int,
+ lr: float,
+ ) -> tuple[float, float]:
+ # data
+ positions, targets = self_play_records.make_training_data(self, td_lambda)
+ n = len(positions)
+ perm = torch.randperm(n, device=positions.device)
+ n_val = max(1, int(n * val_frac))
+ val_idx = perm[:n_val]
+ train_idx = perm[n_val:]
+
+ val_positions = positions[val_idx]
+ val_targets = targets[val_idx]
+ train_positions = positions[train_idx]
+ train_targets = targets[train_idx]
+
+ # training widgets
+ optimizer = torch.optim.Adam(self.parameters(), lr=lr)
+ loss_fn = MSELoss()
+
+ best_corr = -1.0
+ best_state: dict[str, Tensor] | None = None
+ best_mse = float("nan")
+
+ pbar = tqdm(range(epochs), desc="train", unit="ep")
+ for epoch in pbar:
+ self.train()
+ order = torch.randperm(len(train_positions), device=positions.device)
+ for start in range(0, len(train_positions), batch):
+ batch_idx = order[start : start + batch]
+ optimizer.zero_grad()
+ predictions = self(train_positions[batch_idx])
+ loss = loss_fn(predictions, train_targets[batch_idx])
+ loss.backward()
+ optimizer.step()
+
+ with torch.no_grad():
+ self.eval()
+ predictions = self(val_positions)
+ mse = loss_fn(predictions, val_targets).item()
+ corr = torch.corrcoef(torch.stack([predictions, val_targets]))[
+ 0, 1
+ ].item()
+
+ pbar.set_postfix(val_corr=f"{corr:.4f}", val_mse=f"{mse:.4f}")
+
+ if corr > best_corr:
+ best_corr = corr
+ best_state = {
+ name: param.clone() for name, param in self.state_dict().items()
+ }
+ best_mse = mse
+
+ if best_state is not None:
+ self.load_state_dict(best_state)
+
+ return best_corr, best_mse
+
+ def to_c(
+ self,
+ size: int,
+ path: Path,
+ performance: tuple[str, ...] = (),
+ cand: bool = False,
+ ) -> None:
+ prefix = ("cand_" if cand else "") + ("five" if size == 5 else "six")
+ upper = "FIVE" if size == 5 else "SIX"
+
+ W = self.hidden.weight.detach().cpu().numpy()
+ B = self.hidden.bias.detach().cpu().numpy()
+ LW = self.output.weight.detach().cpu().numpy().ravel()
+ 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")
+ for p in performance:
+ _ = f.write(f"/* {p} */\n")
+ _ = f.write(f'#include "weights_{size}.h"\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")
+
+
+# ===================================================================
+# Process helpers
+# ===================================================================
+
+
+def fan_out(cmd: list[str], n: int, procs: int, tag: str, root: Path) -> Iterator[str]:
+ with concurrent.futures.ThreadPoolExecutor(max_workers=procs) as pool:
+ futs = [
+ pool.submit(
+ subprocess.run,
+ cmd,
+ cwd=root,
+ check=True,
+ capture_output=True,
+ text=True,
+ )
+ for _ in range(n)
+ ]
+ for fut in tqdm(
+ concurrent.futures.as_completed(futs), total=n, desc=tag, unit="game"
+ ):
+ yield fut.result().stdout
+
+
+def cand_copy_for(size: int, root: Path) -> None:
+ src = root / "include" / f"weights_{size}.c"
+ if not src.is_file():
+ raise SystemExit(f"{src} not found; build a size-{size} champion first")
+ prefix = "five" if size == 5 else "six"
+ txt = src.read_text().replace(f"{prefix}_", f"cand_{prefix}_")
+ _ = (root / "include" / f"cand_weights_{size}.c").write_text(txt)
+
+
+# ===================================================================
+# Tournament
+# ===================================================================
+
+
+@dataclass(frozen=True)
+class TournamentResult:
+ champ: int
+ cand: int
+ draws: int
+ cand_rate: float
+
+
+def run_matches(
+ *,
+ size: int,
+ depth: int,
+ tournament_eps: float,
+ games: int,
+ procs: int,
+ root: Path,
+) -> TournamentResult:
+ outs = list(
+ fan_out(
+ ["./match", str(size), f"{tournament_eps:.4f}", str(depth)],
+ games,
+ procs,
+ "match",
+ root,
+ )
+ )
+ winners = [line.split()[0] for line in outs if line.strip()]
+ champ = winners.count("current")
+ cand = winners.count("candidate")
+ draws = winners.count("draw")
+ aborts = len(winners) - champ - cand - draws
+ played = len(winners) - aborts
+ return TournamentResult(
+ champ=champ,
+ cand=cand,
+ draws=draws,
+ cand_rate=(cand + 0.5 * draws) / played if played else 0.0,
+ )
+
+
+# ===================================================================
+# The bootstrap loop
+# ===================================================================
+
+
+def setup(size: int, root: Path) -> Net:
+ inp_dim, hidden = NET_DIMS[size]
+ champ = Net(inp_dim, hidden)
+ champ.to_c(
+ size,
+ root / "include" / f"weights_{size}.c",
+ ("v000, random init",),
+ )
+ cand_copy_for(6 if size == 5 else 5, root)
+ return champ
+
+
+def run_step(
+ *,
+ iter: int,
+ size: int,
+ eps_min: float,
+ champ: Net,
+ eps_start: float,
+ eps_decay: float,
+ depth: int,
+ min_ply: int,
+ games: int,
+ procs: int,
+ td_lambda: float,
+ val_frac: float,
+ epochs: int,
+ batch: int,
+ lr: float,
+ tournament_eps: float,
+ tournament_games: int,
+ promote_rate: float,
+ root: Path,
+) -> Net:
+ print(f"\n== iteration {iter} ==")
+
+ # 1. self-play with the installed champion (include/weights_<size>.c)
+ _ = subprocess.run(["make", "selfplay"], cwd=root, check=True, capture_output=True)
+
+ # 2. generate self-play data (one game per process, in parallel)
+ eps = max(eps_min, eps_start * (eps_decay ** (iter - 1)))
+ csv = root / "data" / f"selfplay-{size}-v{iter:03d}.csv"
+ print(f"selfplay, eps={eps:.4f}, depth {depth}, {games} games on {procs} procs")
+ out_rows: list[str] = []
+ for out in fan_out(
+ ["./selfplay", str(size), f"{eps:.4f}", str(depth), str(min_ply)],
+ games,
+ procs,
+ "selfplay",
+ root,
+ ):
+ out_rows.extend(out.splitlines())
+ with open(csv, "w") as f:
+ _ = f.write("\n".join(out_rows))
+ if out_rows:
+ _ = f.write("\n")
+ rows = len(out_rows)
+ print(f" {rows} rows")
+
+ # 3. fine-tune the champion on TD(lambda) targets -> candidate
+ device = next(champ.parameters()).device
+ records = SelfPlayRecords.load_selfplay(csv, size, device)
+ cand = copy.deepcopy(champ)
-# Output of weights
-from tensorflow import transpose
+ corr, mse = cand.fit(
+ self_play_records=records,
+ td_lambda=td_lambda,
+ val_frac=val_frac,
+ epochs=epochs,
+ batch=batch,
+ lr=lr,
+ )
+ cand.to_c(
+ size,
+ root / "include" / f"cand_weights_{size}.c",
+ (f"v{iter:03d}, val_corr={corr:.4f}, val_mse={mse:.4f}",),
+ cand=True,
+ )
-# Building models
-from tensorflow.keras.models import Model
-from tensorflow.keras.layers import Input, Dense
+ # 4. build the match driver (champion + candidate linked)
+ _ = subprocess.run(["make", "match"], cwd=root, check=True, capture_output=True)
+ # 5. tournament + decision
+ tournament_result = run_matches(
+ root=root,
+ size=size,
+ depth=depth,
+ tournament_eps=tournament_eps,
+ games=tournament_games,
+ procs=procs,
+ )
+ decided = tournament_result.champ + tournament_result.cand
+ rate = tournament_result.cand / decided if decided else 0.0
+ print(
+ f"champ {tournament_result.champ} - {tournament_result.cand} cand, {tournament_result.draws} draws, "
+ + f"winrate={rate:.3f}, cand_rate={tournament_result.cand_rate:.3f}"
+ )
+ promote = decided >= int(0.25 * tournament_games) and rate >= promote_rate
-def load_data(size):
- output_len = 2
- # Load training data
- tr_fn = "data/training-"+str(size)+".csv"
- training_csv = pd.read_csv(tr_fn)
- training_data = training_csv
- training_input = np.array(training_data.iloc[:, 0:-output_len])
- training_outcome = np.array(training_data.iloc[:, -output_len:])
- # Load validation data
- val_fn = "data/validation-"+str(size)+".csv"
- val_data = pd.read_csv(val_fn)
- val_input = np.array(val_data.iloc[:, 0:-output_len])
- val_outcome = np.array(val_data.iloc[:, -output_len:])
- return ((training_input, training_outcome), (val_input, val_outcome))
+ if promote:
+ cand.to_c(
+ size,
+ root / "include" / f"weights_{size}.c",
+ (f"v{iter:03d}, val_corr={corr:.4f}",),
+ )
+ print(f"== candidate v{iter:03d} PROMOTED (include/weights_{size}.c) ==")
+ return cand
+ print(f"== candidate v{iter:03d} rejected ==")
+ return champ
-def train(size, model, data, iterations=1, epochs=10, batch=None):
- (tra_input, tra_outcome), (val_input, val_outcome) = data
- results = []
- for i in range(0, iterations):
- print("Iteration {0}/{1}".format(i+1, iterations))
- model.fit(tra_input, tra_outcome, epochs=epochs,
- validation_data=(val_input, val_outcome),
- verbose=True, batch_size=batch)
- val_res = model.evaluate(val_input, val_outcome, verbose=False)
- tra_res = model.evaluate(tra_input, tra_outcome, verbose=False)
- results.append((tra_res, val_res))
- print(val_res)
- write_weights(size, model, str(i+1).zfill(len(str(iterations))), (tra_res, val_res))
- print("\nScores")
- for i, data in enumerate(results):
- print(f"Iteration {i+1}: {data}")
- return results
+# ===================================================================
+# Main
+# ===================================================================
-def make_model(size, magic):
- inputs = Input(shape=(size * size + 3,))
- model = inputs
- model = Dense(magic, activation="relu", use_bias=True)(model)
- model = Dense(2, activation="relu", use_bias=True)(model)
- model = Model(inputs=inputs, outputs=model)
- model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])
- model.summary()
- return model
+def main() -> None:
+ root = Path(__file__).resolve().parent.parent
+ p = argparse.ArgumentParser(
+ description="the TD(lambda) self-play bootstrap loop",
+ 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=10, help="max loop iterations")
+ _ = p.add_argument(
+ "--games", type=int, default=2000, help="selfplay games per iteration"
+ )
+ _ = p.add_argument(
+ "--eps-start",
+ type=float,
+ default=0.15,
+ help="initial self-play epsilon (random-move probability); kept close "
+ "to --tour-eps so the value net trains on the policy it is deployed "
+ "under",
+ )
+ _ = p.add_argument("--eps-min", type=float, default=0.05, help="epsilon floor")
+ _ = p.add_argument(
+ "--eps-decay", type=float, default=0.8, help="per-iteration epsilon decay"
+ )
+ _ = p.add_argument(
+ "--depth",
+ type=int,
+ default=4,
+ help="search depth for selfplay AND the tournament",
+ )
+ _ = p.add_argument(
+ "--min-ply", type=int, default=8, help="record positions at this ply or later"
+ )
+ _ = p.add_argument("--td-lambda", type=float, default=0.6, help="TD(lambda) lambda")
+ _ = p.add_argument(
+ "--epochs", type=int, default=20, help="candidate training epochs"
+ )
+ _ = p.add_argument("--lr", type=float, default=1e-3, help="learning rate")
+ _ = p.add_argument("--batch", type=int, default=128, help="training batch size")
+ _ = p.add_argument(
+ "--val-frac",
+ type=float,
+ default=0.05,
+ help="fraction of rows held out for validation",
+ )
+ _ = p.add_argument(
+ "--tournament-eps",
+ type=float,
+ default=0.05,
+ help="tournament epsilon (random-move probability)",
+ )
+ _ = p.add_argument(
+ "--tournament-games",
+ type=int,
+ default=None,
+ help="tournament games (10/tour_eps if unset)",
+ )
+ _ = p.add_argument(
+ "--promote",
+ type=float,
+ default=0.57,
+ help="candidate win rate among non-drawn tournament games needed to promote (~p<0.03 one-sided binomial at the default 200 games)",
+ )
+ _ = p.add_argument(
+ "--procs",
+ type=int,
+ default=None,
+ help="parallel game processes (half the CPUs if unset)",
+ )
+ args = p.parse_args()
+ if args.tournament_games is None:
+ args.tournament_games = max(20, int(round(10 / args.tournament_eps)))
+ if args.procs is None:
+ args.procs = max(1, (os.cpu_count() or 1) // 2)
-def write_weights(size, model, iteration, performance):
- def fix(val):
- string = str(np.array(val).tolist())
- string = string.replace("[", "{").replace("]", "}")
- return string
- size_str = "five" if size==5 else "six"
- dense_weights = transpose(model.trainable_variables[0], perm=[1, 0])
- dense_biases = model.trainable_variables[1]
- output_weights = transpose(model.trainable_variables[2], perm=[1, 0])
- output_bias = model.trainable_variables[3]
- # Prepare output
- to_output = [(f"{size_str}_dense_weights[{size_str.upper()}_DENSE_NUM][{size_str.upper()}_INP_NUM]", dense_weights),
- (f"{size_str}_dense_biases[{size_str.upper()}_DENSE_NUM]", dense_biases),
- (f"{size_str}_output_weights[2][{size_str.upper()}_DENSE_NUM]", output_weights),
- (f"{size_str}_output_bias[2]", output_bias)]
- # Write to file
- f = open(f"{size_str}_weights-"+iteration+".txt", "w")
- f.write("/*\n")
- model.summary(print_fn=lambda l: f.write(" * "+l+"\n"))
- f.write(" * "+str(performance)+"\n*/\n\n")
- f.write(f"#include \"weights_{size}.h\"\n\n")
- for (name, val) in to_output:
- f.write("const float "+name+" =\n"+fix(val)+";\n\n")
- f.close()
+ (root / "data").mkdir(exist_ok=True)
+ champ = setup(size=args.size, root=root)
+ for iter in range(1, args.iters + 1):
+ champ = run_step(
+ iter=iter,
+ size=args.size,
+ eps_min=args.eps_min,
+ champ=champ,
+ eps_start=args.eps_start,
+ eps_decay=args.eps_decay,
+ depth=args.depth,
+ min_ply=args.min_ply,
+ games=args.games,
+ procs=args.procs,
+ td_lambda=args.td_lambda,
+ val_frac=args.val_frac,
+ epochs=args.epochs,
+ batch=args.batch,
+ lr=args.lr,
+ tournament_eps=args.tournament_eps,
+ tournament_games=args.tournament_games,
+ promote_rate=args.promote,
+ root=root,
+ )
+ print(f"\n== done, champion in include/weights_{args.size}.c ==")
-for size in (5,6):
- data = load_data(size)
- model = make_model(size, 64)
- print("Before training", model.evaluate(data[1][0], data[1][1], verbose=False, batch_size=16))
- results = train(size, model, data, iterations=5, epochs=10, batch=128)
+if __name__ == "__main__":
+ main()
generated by cgit v1.3.1 (git 2.54.0) at 2026-10-15 02:51:49 +0000