From 61e5b8949b2f5e8973c58f0dcc7c20c4bdade029 Mon Sep 17 00:00:00 2001 From: tslil Date: Fri, 25 Sep 2026 14:46:31 +0100 Subject: factor out nn_train.py into a separate python project to make it more manageable, add plotting and offline strength fit --- src/nn_train.py | 885 -------------------------------------------------------- 1 file changed, 885 deletions(-) delete mode 100755 src/nn_train.py (limited to 'src/nn_train.py') diff --git a/src/nn_train.py b/src/nn_train.py deleted file mode 100755 index ae78e43..0000000 --- a/src/nn_train.py +++ /dev/null @@ -1,885 +0,0 @@ -#!/usr/bin/env -S uv run --script -# -# /// script -# requires-python = ">=3.14" -# dependencies = [ -# "numpy>=2.5.3", -# "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 -# 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 . - -import argparse -import concurrent.futures -import json -import math -import os -import random -import subprocess -from collections.abc import Iterator -from dataclasses import dataclass -from itertools import product -from pathlib import Path - -import numpy as np -import torch -from torch import Tensor -from torch.nn import Embedding, Linear, Module, MSELoss -from tqdm import tqdm - -# =================================================================== -# Net dimensions -# =================================================================== - -# (input dim, hidden dim) per board size; must match include/weights_.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 -} - - -# =================================================================== -# 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) - - -# =================================================================== -# Config -# =================================================================== - - -@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 - k: float - - -@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" - elo: float = 1500.0 - rated_games: int = 0 - - 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 - - -# =================================================================== -# Self-play records: games (for ratings) and chains (for training) -# =================================================================== - - -@dataclass(frozen=True) -class Games: - 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} - newest: np.ndarray # [G] game comes from the last loaded file - - def head_to_head(self) -> Iterator[tuple[int, int, float]]: - for i in np.flatnonzero(self.newest & (self.black != self.white)): - black, white, result = ( - int(self.black[i]), - int(self.white[i]), - float(self.result[i]), - ) - # elo score: 1 win, 0.5 draw, 0 loss, from a's perspective - yield black, white, (result + 1.0) / 2.0 - - -@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], - ) - - -# =================================================================== -# Reading game records -# =================================================================== - -_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 -) -> tuple[list[Pool], Games]: - 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, newest = [], [], [], [] - - 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"]) - b, w = _game_colours(game) - black.append(b) - white.append(w) - result.append(game_result) - newest.append(is_newest) - - 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(b == w) - - 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 = Games( - black=np.array(black, dtype=np.int16), - white=np.array(white, dtype=np.int16), - result=np.array(result, dtype=np.float32), - newest=np.array(newest, dtype=np.bool_), - ) - return [pool_of(c) for c in range(n_candidates)], games - - -# =================================================================== -# Training data: fixed rows, refreshed targets -# =================================================================== - - -@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], - ) - - -# =================================================================== -# Network -# =================================================================== - - -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) - - -# =================================================================== -# Process helpers -# =================================================================== - - -def fan_out( - cmds: list[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 cmd in cmds - ] - for fut in tqdm( - concurrent.futures.as_completed(futs), - total=n * len(cmds), - desc=tag, - unit="game", - ): - yield fut.result().stdout - - -def cand_copy_for(size: int, id: 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") - upper = "FIVE" if size == 5 else "SIX" - txt = src.read_text() - for sym in ( - "embedding", - "dense_weights", - "dense_biases", - "logit_weights", - "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) - - -# =================================================================== -# The bootstrap loop -# =================================================================== - - -def update_elo(a: Candidate, b: Candidate, score: float, k: float) -> None: - # score is a's result as an elo score: 1.0 win, 0.5 draw, 0.0 loss - ea = 1.0 / (1.0 + 10.0 ** ((b.elo - a.elo) / 400.0)) - delta = k * (score - ea) - a.elo += delta - b.elo -= delta - - -def run_step( - *, - candidates: list[Candidate], - data_gen_config: DataGenConfig, - 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. rate every head-to-head game played in this iteration - 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, games = load_window(data_files, data_gen_config.size, num_candidates) - print("candidates have", [f"{p.n_rows:,}" for p in pools], "rows.") - - fresh = list(games.head_to_head()) - random.shuffle(fresh) - for a, b, score in fresh: - update_elo(candidates[a], candidates[b], score, data_gen_config.k) - candidates[a].rated_games += 1 - candidates[b].rated_games += 1 - - total_rated_games = sum(c.rated_games for c in candidates) - print( - f"Tournament standings ({total_rated_games:,} rated games, {total_rated_games // len(candidates):,} games/candidate):" - ) - for idx, cand in sorted( - enumerate(candidates), key=lambda t: t[1].elo, reverse=True - ): - print(f"\tcandidate {idx:03d}: Elo {cand.elo:4.1f} Config {cand}") - - # 4. write champ, unconditionally - idx, champ = max(enumerate(candidates), key=lambda t: t[1].elo) - path = data_gen_config.root / "include" / f"weights_{data_gen_config.size}.c" - champ.net.to_c( - size=data_gen_config.size, - ident=None, - path=path, - note=f"{iter=}, elo={champ.elo:.1f}, rated_games={champ.rated_games:,}", - ) - _ = 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}") - - # 5. 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}", - ) - - -# =================================================================== -# Main -# =================================================================== - - -def main() -> None: - root = Path(__file__).resolve().parent.parent - 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="grid points for number of previous data sources to retain", - ) - _ = 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)", - ) - _ = p.add_argument( - "--k-factor", - type=float, - default=32.0, - help="elo k-factor", - ) - - 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, - k=args.k_factor, - ) - 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:,}") - - for iter in range(1, args.iters + 1): - run_step( - candidates=candidates, - data_gen_config=data_gen, - iter=iter, - ) - print(f"\n== done, champion in include/weights_{args.size}.c ==") - - -if __name__ == "__main__": - main() -- cgit v1.3.1