diff options
Diffstat (limited to 'train/src/ctak_train/data.py')
| -rw-r--r-- | train/src/ctak_train/data.py | 300 |
1 files changed, 300 insertions, 0 deletions
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], + ) |