diff options
Diffstat (limited to 'train/src/ctak_train')
| -rw-r--r-- | train/src/ctak_train/__init__.py | 0 | ||||
| -rw-r--r-- | train/src/ctak_train/candidate.py | 125 | ||||
| -rw-r--r-- | train/src/ctak_train/data.py | 300 | ||||
| -rw-r--r-- | train/src/ctak_train/net.py | 129 | ||||
| -rw-r--r-- | train/src/ctak_train/nn_train.py | 299 | ||||
| -rw-r--r-- | train/src/ctak_train/strength.py | 118 |
6 files changed, 971 insertions, 0 deletions
diff --git a/train/src/ctak_train/__init__.py b/train/src/ctak_train/__init__.py new file mode 100644 index 0000000..e69de29 --- /dev/null +++ b/train/src/ctak_train/__init__.py diff --git a/train/src/ctak_train/candidate.py b/train/src/ctak_train/candidate.py new file mode 100644 index 0000000..aa2a6a6 --- /dev/null +++ b/train/src/ctak_train/candidate.py @@ -0,0 +1,125 @@ +# +# 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 math +from dataclasses import dataclass +from itertools import product +from pathlib import Path + +import torch +from torch.nn import MSELoss +from tqdm import tqdm + +from ctak_train.data import DataSplit, Pool +from ctak_train.net import NET_DIMS, Net +from ctak_train.selfplay import cand_copy_for + + +@dataclass(frozen=True) +class DataGenConfig: + size: int + depth: int + min_ply: int + games: int + procs: int + replay_iters: int + root: Path + train_only_self_play: bool + + +@dataclass(frozen=True) +class TrainingGridSearchConfig: + smooth_factor: list[float] + epochs: list[int] + batch: list[int] + lr: list[float] + val_frac: list[float] + + def make_configs(self, data_gen_config: DataGenConfig) -> list["Candidate"]: + inp_dim, hidden = NET_DIMS[data_gen_config.size] + flds = self.__dataclass_fields__.keys() + almost_cands = [ + dict(zip(flds, vals)) for vals in product(*(getattr(self, k) for k in flds)) + ] + for ident, almost_cand in enumerate(almost_cands): + net = Net(inp_dim, hidden) + net.to_c( + size=data_gen_config.size, + ident=ident, + path=data_gen_config.root + / "include" + / f"cand_{ident}_weights_{data_gen_config.size}.c", + note=f"{ident=}, v000, random init", + ) + cand_copy_for( + 6 if data_gen_config.size == 5 else 5, ident, data_gen_config.root + ) + almost_cand["net"] = net + almost_cand["train_only_self_play"] = data_gen_config.train_only_self_play + return [Candidate(**d) for d in almost_cands] + + +@dataclass +class Candidate: + train_only_self_play: bool + smooth_factor: float + epochs: int + batch: int + lr: float + val_frac: float + net: "Net" + + def __str__(self) -> str: + return f"(λ={self.smooth_factor}, E={self.epochs}, BS={self.batch}, η={self.lr:.2e}, V%={self.val_frac * 100.0:.1f})" + + def training_pool(self, pool: Pool) -> Pool: + if self.train_only_self_play: + return pool.select_chains(pool.chain_self_play) + return pool + + def fit(self, pool: Pool) -> tuple[float, float]: + # The train/val split is fixed but targets are refreshed + split = DataSplit.build(self.training_pool(pool), self.val_frac) + if split.n == 0: + return math.nan, math.nan + + optimizer = torch.optim.AdamW(self.net.parameters(), lr=self.lr) + loss_fn = MSELoss() + + pbar = tqdm(range(self.epochs), desc="train", unit="ep") + for _ in pbar: + split.refresh_targets(self.net, self.smooth_factor) + + self.net.train() + for head, cells, parity, target in split.batches(self.batch): + optimizer.zero_grad() + predictions = self.net(head, cells) * parity + loss = loss_fn(predictions, target) + loss.backward() + optimizer.step() + + with torch.no_grad(): + self.net.eval() + val = split.val_idx + predictions = self.net(split.pool.head[val], split.pool.cells[val]) + predictions = predictions * split.pool.parity[val] + targets = split.targets[val] + mse = loss_fn(predictions, targets).item() + corr = torch.corrcoef(torch.stack([predictions, targets]))[0, 1].item() + + pbar.set_postfix(val_corr=f"{corr:.4f}", val_mse=f"{mse:.4f}") + + return corr, mse diff --git a/train/src/ctak_train/data.py b/train/src/ctak_train/data.py new file mode 100644 index 0000000..6dc9ec0 --- /dev/null +++ b/train/src/ctak_train/data.py @@ -0,0 +1,300 @@ +# +# 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 json +from collections.abc import Iterator +from dataclasses import dataclass +from pathlib import Path + +import numpy as np +import torch +from ctak_train.net import NET_DIMS +from torch import Tensor +from torch.nn import Module + + +@dataclass +class GameLog: + iteration: np.ndarray # [G] iteration the game was played in + black: np.ndarray # [G] candidate id playing black + white: np.ndarray # [G] candidate id playing white + result: np.ndarray # [G] outcome from black's perspective, {-1, 0, 1} + + @classmethod + def empty(cls) -> GameLog: + return cls( + iteration=np.empty(0, dtype=np.int32), + black=np.empty(0, dtype=np.int16), + white=np.empty(0, dtype=np.int16), + result=np.empty(0, dtype=np.float32), + ) + + def append(self, newer: GameLog) -> None: + self.iteration = np.concatenate([self.iteration, newer.iteration]) + self.black = np.concatenate([self.black, newer.black]) + self.white = np.concatenate([self.white, newer.white]) + self.result = np.concatenate([self.result, newer.result]) + + +@dataclass(frozen=True) +class Pool: + head: Tensor # [n, 2] fixed head inputs of the root state s_t + cells: Tensor # [n, n_cells] cell states 0..6 of s_t + leaf_head: Tensor # [n, 2] head inputs of the PV leaf pv_t + leaf_cells: Tensor # [n, n_cells] cell states of pv_t + parity: Tensor # [n] +1 when black is to move at s_t + leaf_parity: Tensor # [n] +1 when black is to move at pv_t + terminal: Tensor # [n] leaf is an exact terminal + terminal_value: Tensor # [n] exact terminal value, black's perspective + chain_lens: Tensor # [C] rows per chain + chain_result: Tensor # [C] game result, black's perspective + chain_self_play: Tensor # [C] chain comes from a self vs self game + + @classmethod + def empty(cls, n_cells: int) -> Pool: + def rows(dtype: torch.dtype, width: int = 0) -> Tensor: + return torch.empty((0, width) if width else (0,), dtype=dtype) + + return cls( + head=rows(torch.float32, 2), + cells=rows(torch.uint8, n_cells), + leaf_head=rows(torch.float32, 2), + leaf_cells=rows(torch.uint8, n_cells), + parity=rows(torch.float32), + leaf_parity=rows(torch.float32), + terminal=rows(torch.bool), + terminal_value=rows(torch.float32), + chain_lens=rows(torch.int64), + chain_result=rows(torch.float32), + chain_self_play=rows(torch.bool), + ) + + @property + def n_rows(self) -> int: + return self.head.shape[0] + + def select_chains(self, keep: Tensor) -> Pool: + keep = keep.bool() + rows = torch.repeat_interleave( + torch.arange(self.chain_lens.shape[0], dtype=torch.int64), self.chain_lens + ) + idx = torch.nonzero(keep[rows]).flatten() + return Pool( + head=self.head[idx], + cells=self.cells[idx], + leaf_head=self.leaf_head[idx], + leaf_cells=self.leaf_cells[idx], + parity=self.parity[idx], + leaf_parity=self.leaf_parity[idx], + terminal=self.terminal[idx], + terminal_value=self.terminal_value[idx], + chain_lens=self.chain_lens[keep], + chain_result=self.chain_result[keep], + chain_self_play=self.chain_self_play[keep], + ) + + +_ROW_FIELDS = ( + "head", + "cells", + "leaf_head", + "leaf_cells", + "parity", + "leaf_parity", + "terminal", + "terminal_value", +) + + +def _chain_arrays( + rows: list[dict], n_cells: int, source: Path +) -> dict[str, np.ndarray]: + + def column(part: str, key: str, dtype: np.dtype) -> np.ndarray: + return np.array([r[part][key] for r in rows], dtype=dtype) + + out = { + "head": column("root", "head", np.float32), + "cells": column("root", "cells", np.uint8), + "leaf_head": column("leaf", "head", np.float32), + "leaf_cells": column("leaf", "cells", np.uint8), + "parity": column("root", "parity", np.float32), + "leaf_parity": column("leaf", "parity", np.float32), + "terminal": column("leaf", "terminal", np.bool_), + "terminal_value": column("leaf", "terminal_value", np.float32), + } + for name in ("cells", "leaf_cells"): + if out[name].shape[1] != n_cells: + raise ValueError( + f"{source}: {name} width {out[name].shape[1]}, expected {n_cells}" + ) + return out + + +def _game_colours(game: dict) -> tuple[int, int]: + a, b = int(game["cand_A"]), int(game["cand_B"]) + for r in game["rows"]: + if r["root"]["parity"] > 0: + cand = int(r["cand"]) + return cand, b if cand == a else a + recorded = {int(r["cand"]) for r in game["rows"]} + return (a, b) if a not in recorded else (b, a) + + +def load_window( + files: list[Path], size: int, n_candidates: int, iteration: int +) -> tuple[list[Pool], GameLog]: + """Training pools for every file in the window, plus the game records of + the newest file only - the caller's GameLog already holds everything older.""" + n_cells = NET_DIMS[size][0] - 2 + blocks = [{f: [] for f in _ROW_FIELDS} for _ in range(n_candidates)] + chain_lens: list[list[int]] = [[] for _ in range(n_candidates)] + chain_result: list[list[float]] = [[] for _ in range(n_candidates)] + chain_self_play: list[list[bool]] = [[] for _ in range(n_candidates)] + black, white, result = [], [], [] + + for file_index, path in enumerate(files): + is_newest = file_index == len(files) - 1 + with open(path) as f: + for line in f: + if not line.strip(): + continue + game = json.loads(line) + game_result = float(game["result"]) + self_play = int(game["cand_A"]) == int(game["cand_B"]) + if is_newest: + b, w = _game_colours(game) + black.append(b) + white.append(w) + result.append(game_result) + + by_cand: dict[int, list[dict]] = {} + for r in game["rows"]: + by_cand.setdefault(int(r["cand"]), []).append(r) + + for cand, rows in by_cand.items(): + parity = [r["root"]["parity"] for r in rows] + for sign in (1, -1): + chain = [ + r for r, p in zip(rows, parity) if (p > 0) == (sign > 0) + ] + if not chain: + continue + for name, array in _chain_arrays(chain, n_cells, path).items(): + blocks[cand][name].append(array) + chain_lens[cand].append(len(chain)) + chain_result[cand].append(game_result) + chain_self_play[cand].append(self_play) + + def pool_of(cand: int) -> Pool: + if not chain_lens[cand]: + return Pool.empty(n_cells) + arrays = { + name: torch.from_numpy(np.concatenate(blocks[cand][name])) + for name in _ROW_FIELDS + } + return Pool( + **arrays, + chain_lens=torch.tensor(chain_lens[cand], dtype=torch.int64), + chain_result=torch.tensor(chain_result[cand], dtype=torch.float32), + chain_self_play=torch.tensor(chain_self_play[cand], dtype=torch.bool), + ) + + games = GameLog( + iteration=np.full(len(black), iteration, dtype=np.int32), + black=np.array(black, dtype=np.int16), + white=np.array(white, dtype=np.int16), + result=np.array(result, dtype=np.float32), + ) + return [pool_of(c) for c in range(n_candidates)], games + + +@dataclass +class DataSplit: + pool: Pool + train_idx: Tensor + val_idx: Tensor + targets: Tensor # [n] aligned with pool rows, refreshed per epoch + chain_end: Tensor # [C] last row index of each chain + wave_order: Tensor # [n] rows ordered by distance from their chain's end + wave_bounds: Tensor # [d] start of distance d in wave_order + + @property + def n(self) -> int: + return self.pool.n_rows + + @classmethod + def build(cls, pool: Pool, val_frac: float) -> DataSplit: + n = pool.n_rows + if n == 0: + empty = torch.empty(0, dtype=torch.int64) + return cls(pool, empty, empty, torch.empty(0), empty, empty, empty) + + # Smoothing is G[t] = (1-λ)V[t] + λG[t+1] with G at the chain end. + # Chains are contiguous, so every row that is d plies from its chain's + # end can be updated in one vectorised wavefront step per d. + lens = pool.chain_lens + chain_end = torch.cumsum(lens, dim=0) - 1 + chain_of_row = torch.repeat_interleave( + torch.arange(lens.shape[0], dtype=torch.int64), lens + ) + dist = chain_end[chain_of_row] - torch.arange(n, dtype=torch.int64) + wave_order = torch.argsort(dist, stable=True) + wave_bounds = torch.searchsorted( + dist[wave_order], torch.arange(int(dist.max()) + 2, dtype=torch.int64) + ) + + n_val = max(1, int(n * val_frac)) + perm = torch.randperm(n) + return cls( + pool=pool, + train_idx=perm[n_val:], + val_idx=perm[:n_val], + targets=torch.zeros(n), + chain_end=chain_end, + wave_order=wave_order, + wave_bounds=wave_bounds, + ) + + def refresh_targets(self, net: Module, lam: float) -> None: + pool = self.pool + net.eval() + with torch.no_grad(): + # Net(pv_t), black's perspective; exact terminals override the net + leaf = net(pool.leaf_head, pool.leaf_cells) * pool.leaf_parity + leaf = torch.where(pool.terminal, pool.terminal_value, leaf) + + targets = torch.zeros_like(leaf) + targets[self.chain_end] = pool.chain_result + for lo, hi in zip(self.wave_bounds[1:-1], self.wave_bounds[2:]): + idx = self.wave_order[lo:hi] + targets[idx] = (1.0 - lam) * leaf[idx] + lam * targets[idx + 1] + self.targets = targets + + def batches( + self, batch_size: int + ) -> Iterator[tuple[Tensor, Tensor, Tensor, Tensor]]: + order = torch.randperm(self.train_idx.shape[0]) + train_targets = self.targets[self.train_idx] + for start in range(0, order.shape[0], batch_size): + idx = order[start : start + batch_size] + rows = self.train_idx[idx] + yield ( + self.pool.head[rows], + self.pool.cells[rows], + self.pool.parity[rows], + train_targets[idx], + ) diff --git a/train/src/ctak_train/net.py b/train/src/ctak_train/net.py new file mode 100644 index 0000000..cc8bf5c --- /dev/null +++ b/train/src/ctak_train/net.py @@ -0,0 +1,129 @@ +# +# 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/>. + +from pathlib import Path + +import numpy as np +import torch +from torch import Tensor +from torch.nn import Embedding, Linear, Module + +# (input dim, hidden dim) per board size; must match include/weights_<size>.h +NET_DIMS: dict[int, tuple[int, int]] = { + 5: (27, 64), # FIVE_INP_NUM, FIVE_DENSE_NUM + 6: (38, 64), # SIX_INP_NUM, SIX_DENSE_NUM +} + + +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) + + +def fast_tanh(x: Tensor) -> Tensor: + # Must match fast_tanh in include/nn.c: degree-11 fit of tanh on [-3, 3], + # saturating to +-1 outside. + x2 = x * x + poly = x * ( + 0.997136032 + + x2 + * ( + -0.305215124 + + x2 + * ( + 0.0832466371 + + x2 * (-0.0139947924 + x2 * (0.00123106314 + x2 * -4.29748644e-05)) + ) + ) + ) + return torch.where(x > 3.0, 1.0, torch.where(x < -3.0, -1.0, poly)) + + +def fast_silu(x: Tensor) -> Tensor: + # Must match fast_silu in include/nn.c: [4,2] Pade approximant with cutoff. + # The cutoff below is evaluated at x = 0 (instead of x) because torch + # evaluates both branches of where: 1 + 0.1 * x is 0 for x = -10. + xs = torch.where(x < -4.6, 0.0, x) + poly = xs * (0.5 + xs * (0.25 + xs * (0.05 + 0.004167 * xs))) / (1 + 0.1 * xs) + return torch.where(x < -4.6, 0.0, torch.where(x > 4.6, x, poly)) + + +class Net(Module): + def __init__(self, inp_dim: int, hidden: int) -> None: + super().__init__() + # learned cell-state embedding, see nn.c + self.embedding = Embedding(7, 1) + self.hidden = Linear(inp_dim, hidden, bias=True) + self.output = Linear(hidden, 1, bias=True) + + def forward(self, head: Tensor, cells: Tensor) -> Tensor: + # head: [n, 2] fixed head inputs; cells: [n, n_cells] raw cell states + # (0..6) - the only place they become anything other than integers + inp = torch.cat([head, self.embedding(cells.long()).flatten(1)], dim=1) + hidden = fast_silu(self.hidden(inp)) + return fast_tanh(self.output(hidden)).squeeze(-1) + + def to_c( + self, + *, + size: int, + ident: int | None, + path: Path, + note: str, + ) -> None: + upper = "FIVE" if size == 5 else "SIX" + prefix = (f"cand_{ident}_" if ident is not None else "") + upper + guard = f"WEIGHTS_{upper}" + (f"_{ident}" if ident is not None else "") + + E = self.embedding.weight.detach().cpu().numpy() + 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: + 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}_INP_NUM][{upper}_DENSE_NUM] NN_ALIGNED =\n{{\n{_fmt_matrix(W.T)}\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", + "#endif", + ] + for line in lines: + _ = f.write(line) 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() diff --git a/train/src/ctak_train/strength.py b/train/src/ctak_train/strength.py new file mode 100644 index 0000000..2302606 --- /dev/null +++ b/train/src/ctak_train/strength.py @@ -0,0 +1,118 @@ +# +# 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 math +from pathlib import Path + +import matplotlib +import matplotlib.pyplot as plt +import numpy as np +from ctak_train.candidate import Candidate +from ctak_train.data import GameLog +from scipy.optimize import minimize + +matplotlib.use("Agg") + +# Elo points per unit of logit strength +ELO_PER_LOGIT = 400.0 / math.log(10.0) + +# Spread of the zero-mean prior on the ratings, in logit units. Head-to-head +# results alone leave the likelihood unbounded while the data is thin (a pair +# that has never split draws runs off to +-inf). +PRIOR_SD_LOGIT = 1.0 + + +def fit_strengths(log: GameLog, n_candidates: int) -> tuple[np.ndarray, np.ndarray]: + """Bradley-Terry logistic model (draws scored 0.5) fit in natural-logit + units, mean anchored to 0, reported as Elo points around 1500. Returns + (strengths, +-1 sd). + """ + rated = log.black != log.white + black = log.black[rated].astype(np.int64) + white = log.white[rated].astype(np.int64) + if n_candidates < 2 or black.size == 0: + return np.full(n_candidates, 1500.0), np.zeros(n_candidates) + score = (log.result[rated].astype(np.float64) + 1.0) / 2.0 + + # n_candidates - 1 free parameters; the last candidate anchors the mean to 0 + design = np.eye(n_candidates, n_candidates - 1) + design[n_candidates - 1] = -1.0 + + def nll_grad(theta: np.ndarray) -> tuple[float, np.ndarray]: + rating = design @ theta + x = np.clip(rating[black] - rating[white], -350.0, 350.0) + p = 1.0 / (1.0 + np.exp(-x)) + sp = np.maximum(x, 0.0) + np.log1p(np.exp(-np.abs(x))) # softplus(x) + sp_neg = np.maximum(-x, 0.0) + np.log1p(np.exp(-np.abs(x))) + # nll = score*softplus(-x) + (1-score)*softplus(x): finite for all x + nll = float(np.sum(score * sp_neg + (1.0 - score) * sp)) + nll += 0.5 * float(rating @ rating) / PRIOR_SD_LOGIT**2 + + grad_rating = rating / PRIOR_SD_LOGIT**2 + np.add.at(grad_rating, black, p - score) + np.add.at(grad_rating, white, score - p) + return nll, design.T @ grad_rating + + def hessian(theta: np.ndarray) -> np.ndarray: + rating = design @ theta + x = np.clip(rating[black] - rating[white], -350.0, 350.0) + p = 1.0 / (1.0 + np.exp(-x)) + w = p * (1.0 - p) + + diag = np.bincount(black, weights=w, minlength=n_candidates) + np.bincount( + white, weights=w, minlength=n_candidates + ) + off = np.bincount( + black * n_candidates + white, weights=-w, minlength=n_candidates**2 + ).reshape(n_candidates, n_candidates) + hess_rating = ( + off + + np.diag(diag - np.diag(off)) + + np.eye(n_candidates) / PRIOR_SD_LOGIT**2 + ) + return design.T @ hess_rating @ design + + res = minimize(nll_grad, np.zeros(n_candidates - 1), jac=True, method="L-BFGS-B") + if not res.success: + print(f"strength fit did not converge: {res.message}") + cov = design @ np.linalg.pinv(hessian(res.x)) @ design.T + se = np.sqrt(np.maximum(np.diag(cov), 0.0)) * ELO_PER_LOGIT + return 1500.0 + (design @ res.x) * ELO_PER_LOGIT, se + + +def plot_strengths( + history: list[np.ndarray], candidates: list[Candidate], path: Path +) -> None: + fig, ax = plt.subplots(figsize=(10, 6)) + iterations = range(1, len(history) + 1) + for idx, cand in enumerate(candidates): + ax.plot( + iterations, + [row[idx] for row in history], + marker=".", + ms=3, + lw=1.2, + label=f"cand {idx:03d} {cand}", + ) + ax.axhline(1500.0, color="gray", lw=0.5) + ax.set_xscale("log") + ax.set_xlabel("Iteration") + ax.set_ylabel("Strength (fit on all games so far, mean 1500)") + ax.set_title("Candidate strength, refitted on the full history each iteration") + ax.legend(fontsize=8) + fig.tight_layout() + _ = fig.savefig(path, dpi=150) + plt.close(fig) |