diff options
Diffstat (limited to 'src')
| -rwxr-xr-x | src/nn_train.py | 785 | ||||
| -rw-r--r-- | src/selfplay.c | 7 |
2 files changed, 446 insertions, 346 deletions
diff --git a/src/nn_train.py b/src/nn_train.py index 4de0a5a..f82a369 100755 --- a/src/nn_train.py +++ b/src/nn_train.py @@ -33,26 +33,23 @@ # You should have received a copy of the GNU General Public License # along with ctak. If not, see <https://www.gnu.org/licenses/>. -from collections.abc import Iterator -from dataclasses import dataclass -from itertools import product -from math import log, ceil - import argparse import concurrent.futures import json +import math import os -from pathlib import Path +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 tqdm import tqdm - -from torch.nn import Embedding, MSELoss, Linear, Module -from torch.nn.functional import silu, tanh from torch import Tensor - +from torch.nn import Embedding, Linear, Module, MSELoss +from tqdm import tqdm # =================================================================== # Net dimensions @@ -106,6 +103,9 @@ class DataGenConfig: games: int procs: int replay_iters: int + root: Path + train_only_self_play: bool + k: float @dataclass(frozen=True) @@ -116,292 +116,393 @@ class TrainingGridSearchConfig: lr: list[float] val_frac: list[float] - def make_configs(self) -> list[CandidateTrainingConfig]: + def make_configs(self, data_gen_config: DataGenConfig) -> list["Candidate"]: + inp_dim, hidden = NET_DIMS[data_gen_config.size] flds = self.__dataclass_fields__.keys() - return [ - CandidateTrainingConfig(**dict(zip(flds, vals))) - for vals in product(*(getattr(self, k) for k in flds)) + 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(frozen=True) -class CandidateTrainingConfig: +@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}, V%={self.val_frac * 100.0:.1f})" + 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 -# =================================================================== -# Self-play records -# =================================================================== + 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() -def smooth( - leaf_values: Tensor, - result: float, - lam: float, -) -> Tensor: - # Smooth(leaf_values, result, lambda), rooted at the game outcome: - # G[-1] = result, G[t] = (1-lambda)V[t] + lambda G[t+1] - targets = torch.empty_like(leaf_values) - targets[-1] = result + pbar = tqdm(range(self.epochs), desc="train", unit="ep") + for _ in pbar: + split.refresh_targets(self.net, self.smooth_factor) - for ply in range(len(leaf_values) - 2, -1, -1): - targets[ply] = (1.0 - lam) * leaf_values[ply] + lam * targets[ply + 1] + 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() - return targets + 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}") -SPOOKY: int = 4 + return corr, mse -@dataclass -class Performance: - win: int - loss: int - draw: int +# =================================================================== +# Self-play records: games (for ratings) and chains (for training) +# =================================================================== - def rate(self) -> float: - return float(self.win) / float(max(1, self.win + self.loss + self.draw)) - def __str__(self) -> str: - global SPOOKY - return f"+{str(self.win).zfill(SPOOKY)}-{str(self.loss).zfill(SPOOKY)}={str(self.draw).zfill(SPOOKY)}" +@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 DataSplit: - train_positions: Tensor - train_targets: Tensor - train_parity_signs: Tensor - val_positions: Tensor - val_targets: Tensor - val_parity_signs: Tensor +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) - def batches(self, batch_size: int) -> Iterator[tuple[Tensor, Tensor, Tensor]]: - order = torch.randperm( - len(self.train_positions), device=self.train_positions.device + 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), ) - for start in range(0, len(order), batch_size): - idx = order[start : start + batch_size] - yield ( - self.train_positions[idx], - self.train_targets[idx], - self.train_parity_signs[idx], - ) + @property + def n_rows(self) -> int: + return self.head.shape[0] -@dataclass(frozen=True) -class CandidateRecords: - positions: Tensor # [n, inp_dim] net input of the root state s_t - leaf_positions: Tensor # [n, inp_dim] net input of the PV leaf pv_t - game_ids: Tensor # [n] - results: Tensor # [n] game result R, black's perspective - cand_results: Tensor # [n] R from the recorded candidate's perspective - parity_signs: Tensor # [n] root parity - leaf_parity_signs: Tensor # [n] leaf parity - leaf_terminals: Tensor # [n] leaf is an exact terminal - leaf_terminal_values: Tensor # [n] exact terminal value, black's perspective + 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], + ) - def performance(self) -> Performance: - result = Performance(0, 0, 0) - for game_id in torch.unique(self.game_ids): - in_game = self.game_ids == game_id - game_result = self.cand_results[in_game][0].item() - match game_result: - case 1: - result.win += 1 - case -1: - result.loss += 1 - case _: - result.draw += 1 - return result - def make_training_data( - self, model: Module, training: CandidateTrainingConfig - ) -> DataSplit: - model.eval() - with torch.no_grad(): - # Net(pv_t), black's perspective; exact terminals override the net - leaf_values = model(self.leaf_positions) * self.leaf_parity_signs - leaf_values = torch.where( - self.leaf_terminals, self.leaf_terminal_values, leaf_values +# =================================================================== +# 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) - targets = torch.empty_like(leaf_values) - for game_id in torch.unique(self.game_ids): - in_game = self.game_ids == game_id - game_result = self.results[in_game][0].item() - # self-pair games (i,i) record the candidate on every ply, - # so smooth one chain per colour, not across players - for sign in (1.0, -1.0): - idxs = torch.nonzero( - in_game & (self.parity_signs == sign), as_tuple=True - )[0] - if len(idxs) == 0: +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 - targets[idxs] = smooth( - leaf_values[idxs], game_result, training.smooth_factor - ) + 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) - n_val = max(1, int(len(leaf_values) * training.val_frac)) - perm = torch.randperm(len(leaf_values), device=leaf_values.device) - val_idx, train_idx = perm[:n_val], perm[n_val:] + 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) - return DataSplit( - train_positions=self.positions[train_idx], - train_targets=targets[train_idx], - train_parity_signs=self.parity_signs[train_idx], - val_positions=self.positions[val_idx], - val_targets=targets[val_idx], - val_parity_signs=self.parity_signs[val_idx], + 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 + -@dataclass(frozen=True) -class PlayRecords: - positions: list[Tensor] # [candidate, n_i, inp_dim] - leaf_positions: list[Tensor] # [candidate, n_i, inp_dim] - game_ids: list[Tensor] # [candidate, n_i] - results: list[Tensor] # [candidate, n_i] - cand_results: list[Tensor] # [candidate, n_i] - parity_signs: list[Tensor] # [candidate, n_i] - leaf_parity_signs: list[Tensor] # [candidate, n_i] - leaf_terminals: list[Tensor] # [candidate, n_i] - leaf_terminal_values: list[Tensor] # [candidate, n_i] +# =================================================================== +# Training data: fixed rows, refreshed targets +# =================================================================== - def index(self, index: int) -> CandidateRecords: - return CandidateRecords( - positions=self.positions[index], - leaf_positions=self.leaf_positions[index], - game_ids=self.game_ids[index], - results=self.results[index], - cand_results=self.cand_results[index], - parity_signs=self.parity_signs[index], - leaf_parity_signs=self.leaf_parity_signs[index], - leaf_terminals=self.leaf_terminals[index], - leaf_terminal_values=self.leaf_terminal_values[index], - ) - @classmethod - def load_merge(cls, files: list[Path], size: int) -> PlayRecords: - # one JSON record per line: - # {"result": R, "rows": [{"cand", "root": {...}, "leaf": {...}}]} - inp_dim = NET_DIMS[size][0] - n_cells = inp_dim - 2 +@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 - games: list[dict] = [] - for file in files: - with open(file) as f: - for line in f: - if line.strip(): - games.append(json.loads(line)) + @property + def n(self) -> int: + return self.pool.n_rows - def board(row_part: dict) -> list[float]: - return [*row_part["head"], *row_part["cells"]] + @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) - n_cands = max(r["cand"] for g in games for r in g["rows"]) + 1 - fields = ( - "positions", - "leaf_positions", - "game_ids", - "results", - "cand_results", - "parity_signs", - "leaf_parity_signs", - "leaf_terminals", - "leaf_terminal_values", + # 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) ) - per_cand = {c: {k: [] for k in fields} for c in range(n_cands)} - for game_id, game in enumerate(games): - result = float(game["result"]) - for cand in range(n_cands): - rows = [r for r in game["rows"] if r["cand"] == cand] - if not rows: - continue - for part in ("root", "leaf"): - for r in rows: - if len(r[part]["cells"]) != n_cells: - raise ValueError( - f"{files[-1]}: {len(r[part]['cells'])} cells, " - f"expected {n_cells} for size {size}" - ) - parity = np.array([r["root"]["parity"] for r in rows], dtype=np.float32) - per_cand[cand]["positions"].append( - torch.from_numpy( - np.array([board(r["root"]) for r in rows], dtype=np.float32) - ) - ) - per_cand[cand]["leaf_positions"].append( - torch.from_numpy( - np.array([board(r["leaf"]) for r in rows], dtype=np.float32) - ) - ) - per_cand[cand]["game_ids"].append( - torch.full((len(rows),), game_id, dtype=torch.int64) - ) - per_cand[cand]["results"].append( - torch.full((len(rows),), result, dtype=torch.float32) - ) - per_cand[cand]["cand_results"].append(torch.from_numpy(parity * result)) - per_cand[cand]["parity_signs"].append(torch.from_numpy(parity)) - per_cand[cand]["leaf_parity_signs"].append( - torch.from_numpy( - np.array([r["leaf"]["parity"] for r in rows], dtype=np.float32) - ) - ) - per_cand[cand]["leaf_terminals"].append( - torch.from_numpy( - np.array( - [r["leaf"]["terminal"] == 1 for r in rows], - dtype=np.bool_, - ) - ) - ) - per_cand[cand]["leaf_terminal_values"].append( - torch.from_numpy( - np.array( - [r["leaf"]["terminal_value"] for r in rows], - dtype=np.float32, - ) - ) - ) - - def merged(cand: int, k: str, empty: Tensor) -> Tensor: - blocks = per_cand[cand][k] - return torch.concat(blocks, dim=0) if blocks else empty + n_val = max(1, int(n * val_frac)) + perm = torch.randperm(n) return cls( - **{ - k: [ - merged( - cand, - k, - torch.empty( - (0, inp_dim) if k.endswith("positions") else (0,), - dtype=torch.float32 - if k not in ("game_ids", "leaf_terminals") - else (torch.int64 if k == "game_ids" else torch.bool), - ), - ) - for cand in range(n_cands) - ] - for k in fields - } + 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__() @@ -410,56 +511,24 @@ class Net(Module): self.hidden = Linear(inp_dim, hidden, bias=True) self.output = Linear(hidden, 1, bias=True) - def forward(self, x: Tensor) -> Tensor: - # x: [n, 2 + n_cells]; the first 2 columns are the fixed head - # inputs, the rest are raw cell states (0..6) - n = x.shape[0] - inp = torch.cat([x[:, :2], self.embedding(x[:, 2:].long()).view(n, -1)], dim=1) - hidden = silu(self.hidden(inp)) - return tanh(self.output(hidden)).squeeze(-1) - - def fit( - self, *, records: CandidateRecords, config: CandidateTrainingConfig - ) -> tuple[float, float]: - # training widgets - optimizer = torch.optim.AdamW(self.parameters(), lr=config.lr) - loss_fn = MSELoss() - - pbar = tqdm(range(config.epochs), desc="train", unit="ep") - for _ in pbar: - split = records.make_training_data(self, config) - - self.train() - for positions, targets, parity_signs in split.batches(config.batch): - optimizer.zero_grad() - predictions = self(positions) * parity_signs - loss = loss_fn(predictions, targets) - loss.backward() - optimizer.step() - - with torch.no_grad(): - self.eval() - predictions = self(split.val_positions) * split.val_parity_signs - mse = loss_fn(predictions, split.val_targets).item() - corr = torch.corrcoef(torch.stack([predictions, split.val_targets]))[ - 0, 1 - ].item() - - pbar.set_postfix(val_corr=f"{corr:.4f}", val_mse=f"{mse:.4f}") - - return corr, mse + 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, - id: int | None, + ident: int | None, path: Path, note: str, ) -> 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 "") + 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() @@ -478,7 +547,7 @@ class Net(Module): 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", + "#endif", ] for line in lines: _ = f.write(line) @@ -537,30 +606,26 @@ def cand_copy_for(size: int, id: int, root: Path) -> None: # =================================================================== -def setup(num_candidates: int, data_gen: DataGenConfig, root: Path) -> list[Net]: - inp_dim, hidden = NET_DIMS[data_gen.size] - candidates = [Net(inp_dim, hidden) for _ in range(num_candidates)] - for id, cand in enumerate(candidates): - cand.to_c( - size=data_gen.size, - id=id, - path=root / "include" / f"cand_{id}_weights_{data_gen.size}.c", - note=f"{id=}, v000, random init", - ) - cand_copy_for(6 if data_gen.size == 5 else 5, id, root) - return candidates +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[Net], - data_gen: DataGenConfig, + candidates: list[Candidate], + data_gen_config: DataGenConfig, iter: int, - root: Path, - training_configs: list[CandidateTrainingConfig], ): def data_path(i: int) -> Path: - return root / "data" / f"selfplay-{data_gen.size}-v{i:03d}.json" + return ( + data_gen_config.root + / "data" + / f"selfplay-{data_gen_config.size}-v{i:03d}.json" + ) print(f"\n== iteration {iter} ==") @@ -572,30 +637,34 @@ def run_step( # 1. rebuild for i, j in cand_pairs: _ = subprocess.run( - ["make", f"selfplay-{i}-{j}"], cwd=root, check=True, capture_output=True + ["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.games=}, " - f"{data_gen.size=}, {data_gen.depth=}, {data_gen.min_ply=}" + 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.size), - str(data_gen.depth), - str(data_gen.min_ply), + str(data_gen_config.size), + str(data_gen_config.depth), + str(data_gen_config.min_ply), ] for (i, j) in cand_pairs ], - data_gen.games, - data_gen.procs, + data_gen_config.games, + data_gen_config.procs, "selfplay", - root, + data_gen_config.root, ): out_games.extend(out.splitlines()) print(f"writing {len(out_games):,} games to {data}") @@ -604,43 +673,57 @@ def run_step( if out_games: _ = f.write("\n") - # 3. compute performance + # 3. rate every head-to-head game played in this iteration data_files = [ - data_path(i) for i in range(max(iter - data_gen.replay_iters, 1), iter + 1) + 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=" ") - records = PlayRecords.load_merge(data_files, data_gen.size) - print("candidates have", [f"{r.shape[0]:,}" for r in records.results], "rows.") - results = [ - (i, cand, records.index(i).performance()) for i, cand in enumerate(candidates) - ] + pools, games = load_window(data_files, data_gen_config.size, num_candidates) + print("candidates have", [f"{p.n_rows:,}" for p in pools], "rows.") - print("Tournament results:") - for idx, _, perf in sorted(results, key=lambda t: t[2].rate(), reverse=True): - print( - f"\tcandidate {idx:03d} {perf} {perf.rate() * 100.0:.2f}% {training_configs[idx]}" - ) + 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 - # 4. write champ - idx, champ, perf = max(results, key=lambda t: t[2].rate()) - path = root / "include" / f"weights_{data_gen.size}.c" - champ.to_c( - size=data_gen.size, - id=None, + 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=}, perf={str(perf)}, rate={perf.rate() * 100.0:.3f}", + note=f"{iter=}, elo={champ.elo:.1f}, rated_games={champ.rated_games:,}", ) - _ = subprocess.run(["make", "clean"], cwd=root, check=True, capture_output=True) - _ = subprocess.run(["make", "native"], cwd=root, check=True, capture_output=True) - print(f"Wrote candidate {idx:03d} to {str(path)}") + _ = 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(records=records.index(idx), config=training_configs[idx]) - cand.to_c( - size=data_gen.size, - id=idx, - path=root / "include" / f"cand_{idx}_weights_{data_gen.size}.c", + 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}", ) @@ -734,9 +817,25 @@ def main() -> None: 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)) @@ -747,6 +846,9 @@ def main() -> None: 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, @@ -756,23 +858,20 @@ def main() -> None: val_frac=args.val_frac, ) - training_configs = tgsc.make_configs() + candidates = tgsc.make_configs(data_gen) (root / "data").mkdir(exist_ok=True) - candidates = setup( - num_candidates=len(training_configs), data_gen=data_gen, root=root - ) - num_params = sum(p.numel() for p in candidates[0].parameters() if p.requires_grad) + 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=data_gen, + data_gen_config=data_gen, iter=iter, - root=root, - training_configs=training_configs, ) print(f"\n== done, champion in include/weights_{args.size}.c ==") diff --git a/src/selfplay.c b/src/selfplay.c index b2a5821..37f94fa 100644 --- a/src/selfplay.c +++ b/src/selfplay.c @@ -169,7 +169,8 @@ static int play_game(const uint8_t min_ply, const nn_function fn_1, // final outcome of the played game, from black's perspective const int8_t R = win_value(state->won); - fprintf(stdout, "{\"result\":%d,\"rows\":[", R); + fprintf(stdout, "{\"result\":%d,\"cand_A\":%u,\"cand_B\":%u,\"rows\":[", R, + CAND_A, CAND_B); for (uint32_t k = 0; k < n_rows; k++) { fprintf(stdout, "%s{\"cand\":%u,\"root\":{\"parity\":%d,", k ? "," : "", row_cand_id[k], row_mover_sign[k]); @@ -188,8 +189,8 @@ const char *usage = "selfplay, one game of training data between candidates.\n\n\ Usage: selfplay <size> <depth> <min_ply>\n\n\ Plays a single game of <size>x<size>, with <min_ply> moves made randomly. Only the deterministic plies are written to stdout, one JSON record per game, from ply <min_ply> onwards:\n\n\ -\t{\"result\": R, \"rows\": [{\"cand\": id, \"root\": {\"parity\", \"head\", \"cells\"}, \"leaf\": {\"parity\", \"head\", \"cells\", \"terminal\", \"terminal_value\"}}]}\n\n\ -R is the outcome of the game from black's perspective. Each row records the net input of the root state s_t and the PV leaf selected at that ply; the leaf value Net(pv) is taken from the trained net. nn_train.py runs many of these in parallel to generate training data.\n\n\ +\t{\"result\": R, \"cand_A\": i, \"cand_B\": j, \"rows\": [{\"cand\": id, \"root\": {\"parity\", \"head\", \"cells\"}, \"leaf\": {\"parity\", \"head\", \"cells\", \"terminal\", \"terminal_value\"}}]}\n\n\ +R is the outcome of the game from black's perspective; cand_A and cand_B are the ids of the two candidates that played the game. Each row records the net input of the root state s_t and the PV leaf selected at that ply; the leaf value Net(pv) is taken from the trained net. nn_train.py runs many of these in parallel to generate training data.\n\n\ Copyright (C) 2026, tslil clingman\n\n\ 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."; |