diff options
Diffstat (limited to 'src/nn_train.py')
| -rwxr-xr-x | src/nn_train.py | 429 |
1 files changed, 220 insertions, 209 deletions
diff --git a/src/nn_train.py b/src/nn_train.py index 0e08da7..2eb9b0c 100755 --- a/src/nn_train.py +++ b/src/nn_train.py @@ -35,6 +35,7 @@ from collections.abc import Iterator from dataclasses import dataclass +from math import log, ceil import argparse import concurrent.futures @@ -114,67 +115,43 @@ def td_lambda_targets( return targets -@dataclass(frozen=True) -class SelfPlayRecords: - positions: Tensor # [n, inp_dim] - game_ids: Tensor # [n] id of the game the position belongs to - results: Tensor # [n] final played result from black's perspective - terminals: Tensor # [n] true when the recorded position is a terminal leaf - terminal_values: Tensor # [n] exact value of a terminal leaf, otherwise 0 - parity_signs: Tensor # [n] ±1 converting the net's output to "absolute" perspective +@dataclass +class Performance: + win: int + loss: int + draw: int - @classmethod - def load_selfplay( - cls, fn: Path, size: int, device: torch.device - ) -> SelfPlayRecords: - # [net input] [game_id] [played result] [terminal flag] [leaf result] [parity sign] - df = pd.read_csv(fn, header=None) - inp_dim = NET_DIMS[size][0] - return cls( - positions=torch.tensor( - df.iloc[:, 0:inp_dim].to_numpy(dtype=np.float32), - dtype=torch.float32, - device=device, - ), - game_ids=torch.tensor( - df.iloc[:, inp_dim].to_numpy(dtype=np.int64), - device=device, - ), - results=torch.tensor( - df.iloc[:, inp_dim + 1].to_numpy(dtype=np.float32), - dtype=torch.float32, - device=device, - ), - terminals=torch.tensor( - df.iloc[:, inp_dim + 2].to_numpy(dtype=np.bool_), - dtype=torch.bool, - device=device, - ), - terminal_values=torch.tensor( - df.iloc[:, inp_dim + 3].to_numpy(dtype=np.float32), - dtype=torch.float32, - device=device, - ), - parity_signs=torch.tensor( - df.iloc[:, inp_dim + 4].to_numpy(dtype=np.float32), - dtype=torch.float32, - device=device, - ), - ) + def rate(self) -> float: + return float(self.win) / float(max(1, self.win + self.loss + self.draw)) - @classmethod - def load_merge( - cls, files: list[Path], size: int, device: torch.device - ) -> SelfPlayRecords: - sprs = [SelfPlayRecords.load_selfplay(file, size, device) for file in files] - return cls( - positions=torch.concat([spr.positions for spr in sprs]), - game_ids=torch.concat([spr.game_ids for spr in sprs]), - results=torch.concat([spr.results for spr in sprs]), - terminals=torch.concat([spr.terminals for spr in sprs]), - terminal_values=torch.concat([spr.terminal_values for spr in sprs]), - parity_signs=torch.concat([spr.parity_signs for spr in sprs]), - ) + def __str__(self) -> str: + width = ceil(log(max(1, self.win, self.loss, self.draw), 10)) + return f"+{str(self.win).zfill(width)}-{str(self.loss).zfill(width)}={str(self.draw).zfill(width)}" + + +@dataclass(frozen=True) +class CandidateRecords: + positions: list[Tensor] # [n, inp_dim] + game_ids: list[Tensor] # [n] + results: list[Tensor] # [n] + cand_results: list[Tensor] # [n] + terminals: list[Tensor] # [n] + terminal_values: list[Tensor] # [n] + parity_signs: list[Tensor] # [n] + + 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, td_lambda: float @@ -201,6 +178,85 @@ class SelfPlayRecords: return self.positions, targets, self.parity_signs +@dataclass(frozen=True) +class PlayRecords: + 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] + terminals: list[Tensor] # [candidate, n_i] + terminal_values: list[Tensor] # [candidate, n_i] + parity_signs: list[Tensor] # [candidate, n_i] + + def index(self, index: int) -> CandidateRecords: + return CandidateRecords( + positions=self.positions[index], + game_ids=self.game_ids[index], + results=self.results[index], + cand_results=self.cand_results[index], + terminals=self.terminals[index], + terminal_values=self.terminal_values[index], + parity_signs=self.parity_signs[index], + ) + + @classmethod + def load_selfplay(cls, fn, size): + df = pd.read_csv(fn, header=None) + inp_dim = NET_DIMS[size][0] + + fields = { + "parity_signs": df.iloc[:, 0].to_numpy(dtype=np.float32), + "positions": df.iloc[:, 1 : inp_dim + 1].to_numpy(dtype=np.float32), + "results": df.iloc[:, inp_dim + 1].to_numpy(dtype=np.float32), + "terminals": df.iloc[:, inp_dim + 2].to_numpy(dtype=np.bool_), + "terminal_values": df.iloc[:, inp_dim + 3].to_numpy(dtype=np.float32), + "game_ids": df.iloc[:, inp_dim + 4].to_numpy(dtype=np.int64), + "cand_results": df.iloc[:, inp_dim + 6].to_numpy(dtype=np.float32), + } + cand_ids = df.iloc[:, inp_dim + 5].to_numpy(dtype=np.int64) + + n_cands = int(cand_ids.max()) + 1 + gathered = {k: [] for k in fields} + for cid in range(n_cands): + mask = cand_ids == cid + for k, arr in fields.items(): + gathered[k].append(torch.from_numpy(arr[mask])) + + return cls(**gathered) + + @classmethod + def load_merge(cls, files: list[Path], size: int) -> PlayRecords: + sprs = [PlayRecords.load_selfplay(file, size) for file in files] + # for each field, gather [file0[i], file1[i], ...] for every index i + # and concat along dim 0 -> merged list indexed by cand id + return cls( + positions=[ + torch.concat(vals, dim=0) for vals in zip(*(s.positions for s in sprs)) + ], + game_ids=[ + torch.concat(vals, dim=0) for vals in zip(*(s.game_ids for s in sprs)) + ], + results=[ + torch.concat(vals, dim=0) for vals in zip(*(s.results for s in sprs)) + ], + cand_results=[ + torch.concat(vals, dim=0) + for vals in zip(*(s.cand_results for s in sprs)) + ], + terminals=[ + torch.concat(vals, dim=0) for vals in zip(*(s.terminals for s in sprs)) + ], + terminal_values=[ + torch.concat(vals, dim=0) + for vals in zip(*(s.terminal_values for s in sprs)) + ], + parity_signs=[ + torch.concat(vals, dim=0) + for vals in zip(*(s.parity_signs for s in sprs)) + ], + ) + + # =================================================================== # Network # =================================================================== @@ -225,7 +281,7 @@ class Net(Module): def fit( self, *, - self_play_records: SelfPlayRecords, + records: CandidateRecords, td_lambda: float, val_frac: float, epochs: int, @@ -233,9 +289,7 @@ class Net(Module): lr: float, ) -> tuple[float, float]: # data - positions, targets, parity_signs = self_play_records.make_training_data( - self, td_lambda - ) + positions, targets, parity_signs = records.make_training_data(self, td_lambda) n = len(positions) perm = torch.randperm(n, device=positions.device) n_val = max(1, int(n * val_frac)) @@ -258,7 +312,7 @@ class Net(Module): best_mse = float("nan") pbar = tqdm(range(epochs), desc="train", unit="ep") - for epoch in pbar: + for _ in pbar: self.train() order = torch.randperm(len(train_positions), device=positions.device) for start in range(0, len(train_positions), batch): @@ -295,13 +349,14 @@ class Net(Module): def to_c( self, + *, size: int, + id: int | None, path: Path, - performance: tuple[str, ...] = (), - cand: bool = False, + note: str, ) -> None: - prefix = ("cand_" if cand else "") + ("five" if size == 5 else "six") upper = "FIVE" if size == 5 else "SIX" + prefix = (f"cand_{id}_" if id is not None else "") + upper E = self.embedding.weight.detach().cpu().numpy() W = self.hidden.weight.detach().cpu().numpy() @@ -311,8 +366,7 @@ class Net(Module): with open(path, "w") as f: _ = f.write("/* generated by nn_train.py, do not edit by hand */\n") - for p in performance: - _ = f.write(f"/* {p} */\n") + _ = f.write(f"/* {note} */\n") _ = f.write(f'#include "weights_{size}.h"\n\n') _ = f.write( f"const float {prefix}_embedding[{len(E)}] =\n" @@ -338,7 +392,9 @@ class Net(Module): # =================================================================== -def fan_out(cmd: list[str], n: int, procs: int, tag: str, root: Path) -> Iterator[str]: +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( @@ -350,57 +406,32 @@ def fan_out(cmd: list[str], n: int, procs: int, tag: str, root: Path) -> Iterato text=True, ) for _ in range(n) + for cmd in cmds ] for fut in tqdm( - concurrent.futures.as_completed(futs), total=n, desc=tag, unit="game" + concurrent.futures.as_completed(futs), + total=n * len(cmds), + desc=tag, + unit="game", ): yield fut.result().stdout -def cand_copy_for(size: int, root: Path) -> None: +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") - prefix = "five" if size == 5 else "six" - txt = src.read_text().replace(f"{prefix}_", f"cand_{prefix}_") - _ = (root / "include" / f"cand_weights_{size}.c").write_text(txt) - - -# =================================================================== -# Tournament -# =================================================================== - - -@dataclass(frozen=True) -class TournamentResult: - champ: int - cand: int - draws: int - - -def run_matches( - *, - depth: int, - games: int, - procs: int, - root: Path, - size: int, - match_random_opening_length: int, -) -> TournamentResult: - outs = list( - fan_out( - ["./match", str(size), f"{match_random_opening_length}", str(depth)], - games, - procs, - "match", - root, - ) - ) - winners = [line.split()[0] for line in outs if line.strip()] - champ = winners.count("current") - cand = winners.count("candidate") - draws = winners.count("draw") - return TournamentResult(champ=champ, cand=cand, draws=draws) + 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}") + _ = (root / "include" / f"cand_{id}_weights_{size}.c").write_text(txt) # =================================================================== @@ -408,120 +439,115 @@ def run_matches( # =================================================================== -def setup(size: int, root: Path) -> Net: +def setup(num_candidates: int, size: int, root: Path) -> list[Net]: inp_dim, hidden = NET_DIMS[size] - champ = Net(inp_dim, hidden) - champ.to_c( - size, - root / "include" / f"weights_{size}.c", - ("v000, random init",), - ) - cand_copy_for(6 if size == 5 else 5, root) - return champ + candidates = [Net(inp_dim, hidden) for _ in range(num_candidates)] + for id, cand in enumerate(candidates): + cand.to_c( + size=size, + id=id, + path=root / "include" / f"cand_{id}_weights_{size}.c", + note=f"{id=}, v000, random init", + ) + cand_copy_for(6 if size == 5 else 5, id, root) + return candidates def run_step( *, batch: int, - champ: Net, + candidates: list[Net], depth: int, epochs: int, games: int, iter: int, lr: float, - matches: int, - match_random_opening_length: int, min_ply: int, procs: int, - promote_rate: float, replay_iters: int, root: Path, size: int, td_lambda: float, val_frac: float, -) -> Net: - def make_csv(i: int) -> Path: +): + def csv_path(i: int) -> Path: return root / "data" / f"selfplay-{size}-v{i:03d}.csv" print(f"\n== iteration {iter} ==") + num_candidates = len(candidates) + cand_pairs = [ + (i, j) for i in range(num_candidates) for j in range(i + 1, num_candidates) + ] + # 1. rebuild - _ = subprocess.run(["make", "selfplay"], cwd=root, check=True, capture_output=True) - _ = subprocess.run(["make", "native"], cwd=root, check=True, capture_output=True) + for i, j in cand_pairs: + _ = subprocess.run( + ["make", f"selfplay-{i}-{j}"], cwd=root, check=True, capture_output=True + ) # 2. generate self-play data - csv = make_csv(iter) - print(f"selfplay: {size=}, {depth=}, {min_ply=}") + csv = csv_path(iter) + print(f"selfplay: {num_candidates=}, {size=}, {depth=}, {min_ply=}") out_rows: list[str] = [] for out in fan_out( - ["./selfplay", str(size), str(depth), str(min_ply)], + [ + [f"./selfplay-{i}-{j}", str(size), str(depth), str(min_ply)] + for (i, j) in cand_pairs + ], games, procs, "selfplay", root, ): out_rows.extend(out.splitlines()) + print("writing to", csv, "with", len(out_rows)) with open(csv, "w") as f: _ = f.write("\n".join(out_rows)) if out_rows: _ = f.write("\n") - # 3. fine-tune the champion on TD(lambda) targets -> candidate - device = next(champ.parameters()).device - csvs = [make_csv(i) for i in range(max(iter - replay_iters, 1), iter + 1)] + # 3. compute performance + csvs = [csv_path(i) for i in range(max(iter - replay_iters, 1), iter + 1)] print("Loading data from", [str(c) for c in csvs], "...", end=" ") - records = SelfPlayRecords.load_merge(csvs, size, device) - print(f"{records.results.shape[0]:,} rows") - cand = copy.deepcopy(champ) + records = PlayRecords.load_merge(csvs, 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) + ] - corr, mse = cand.fit( - self_play_records=records, - td_lambda=td_lambda, - val_frac=val_frac, - epochs=epochs, - batch=batch, - lr=lr, - ) - cand.to_c( - size, - root / "include" / f"cand_weights_{size}.c", - (f"v{iter:03d}, val_corr={corr:.4f}, val_mse={mse:.4f}",), - cand=True, - ) + for idx, _, perf in sorted(results, key=lambda t: t[2].rate(), reverse=True): + print(f"\t cand {idx:03d} {perf} {perf.rate() * 100.0:.2f}%") - # 4. build the match driver (champion + candidate) - _ = subprocess.run(["make", "match"], cwd=root, check=True, capture_output=True) - - promote = True - rate = -1.0 - if promote_rate > 0: - # 5. tournament + decision - print(f"match: {size=}, {match_random_opening_length=}, {depth=}") - tournament_result = run_matches( - depth=depth, - games=matches, - procs=procs, - root=root, - size=size, - match_random_opening_length=match_random_opening_length, - ) - decided = tournament_result.champ + tournament_result.cand - rate = tournament_result.cand / decided if decided else 0.0 - print( - f"champ {tournament_result.champ} - {tournament_result.cand} cand, {tournament_result.draws} draws, winrate={rate:.3f}" + # 4. fine-tune the champion on TD(lambda) targets -> candidate + for idx, cand in enumerate(candidates): + corr, mse = cand.fit( + records=records.index(idx), + td_lambda=td_lambda, + val_frac=val_frac, + epochs=epochs, + batch=batch, + lr=lr, ) - promote = decided >= int(0.25 * matches) and rate > promote_rate - - if promote: cand.to_c( - size, - root / "include" / f"weights_{size}.c", - (f"v{iter:03d}, val_corr={corr:.4f}, rate={rate:.4f}",), + size=size, + id=idx, + path=root / "include" / f"cand_{idx}_weights_{size}.c", + note=f"id={idx}, v{iter:03d}, val_corr={corr:.4f}, val_mse={mse:.4f}", ) - print(f"== candidate v{iter:03d} PROMOTED (include/weights_{size}.c) ==") - return cand - print(f"== candidate v{iter:03d} rejected ==") - return champ + + # 5. write champ + _, champ, perf = max(results, key=lambda t: t[2].rate()) + champ.to_c( + size=size, + id=None, + path=root / "include" / f"weights_{size}.c", + note=f"{iter=}, perf={str(perf)}", + ) + _ = subprocess.run(["make", "clean"], cwd=root, check=True, capture_output=True) + for i, j in cand_pairs: + (root / f"selfplay-{i}-{j}").unlink(missing_ok=True) + _ = subprocess.run(["make", "native"], cwd=root, check=True, capture_output=True) # =================================================================== @@ -551,14 +577,20 @@ def main() -> None: _ = p.add_argument( "--games", type=int, - default=1000, + default=100, help="selfplay games per iteration", ) _ = p.add_argument( + "--num-candidates", + type=int, + default=5, + help="number of distinct candidates to train", + ) + _ = p.add_argument( "--replay-iters", type=int, default=4, - help="number of previous self-play instances to mix in with current dataset during training", + help="number of previous data sources to retain", ) _ = p.add_argument( "--td-lambda", @@ -591,24 +623,6 @@ def main() -> None: help="fraction of rows held out for validation", ) _ = p.add_argument( - "--matches", - type=int, - default=200, - help="how many matches in a tournament", - ) - _ = p.add_argument( - "--match-random-opening-length", - type=int, - default=2, - help="length of random opening in a match", - ) - _ = p.add_argument( - "--promote", - type=float, - default=0.53, - help="candidate win rate among non-drawn matches needed to promote", - ) - _ = p.add_argument( "--procs", type=int, default=None, @@ -633,25 +647,22 @@ def main() -> None: args.procs = max(1, (os.cpu_count() or 1) // 2) (root / "data").mkdir(exist_ok=True) - champ = setup(size=args.size, root=root) + candidates = setup(num_candidates=args.num_candidates, size=args.size, root=root) - num_params = sum(p.numel() for p in champ.parameters() if p.requires_grad) + num_params = sum(p.numel() for p in candidates[0].parameters() if p.requires_grad) print(f"Number of trainable parameters: {num_params:,}") for iter in range(1, args.iters + 1): - champ = run_step( + run_step( batch=args.batch, - champ=champ, + candidates=candidates, depth=args.depth, epochs=args.epochs, games=args.games, iter=iter, lr=args.lr, - matches=args.matches, - match_random_opening_length=args.match_random_opening_length, min_ply=args.min_ply, procs=args.procs, - promote_rate=args.promote, replay_iters=args.replay_iters, root=root, size=args.size, |