diff options
Diffstat (limited to 'src/nn_train.py')
| -rwxr-xr-x | src/nn_train.py | 98 |
1 files changed, 62 insertions, 36 deletions
diff --git a/src/nn_train.py b/src/nn_train.py index 52549d8..8025b68 100755 --- a/src/nn_train.py +++ b/src/nn_train.py @@ -152,6 +152,18 @@ class SelfPlayRecords: ), ) + @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]), + plies=torch.concat([spr.plies for spr in sprs]), + results=torch.concat([spr.results for spr in sprs]), + ) + def make_training_data( self, model: Module, td_lambda: float ) -> tuple[Tensor, Tensor]: @@ -382,26 +394,30 @@ def setup(size: int, root: Path) -> Net: def run_step( *, - iter: int, - size: int, - eps_min: float, + batch: int, champ: Net, - eps_start: float, - eps_decay: float, depth: int, - min_ply: int, + epochs: int, + eps_decay: float, + eps_min: float, + eps_start: float, games: int, + iter: int, + lr: float, + min_ply: int, procs: int, + promote_rate: float, + replay_iters: int, + root: Path, + size: int, td_lambda: float, - val_frac: float, - epochs: int, - batch: int, - lr: float, tournament_eps: float, tournament_games: int, - promote_rate: float, - root: Path, + val_frac: float, ) -> Net: + def make_csv(i: int) -> Path: + return root / "data" / f"selfplay-{size}-v{i:03d}.csv" + print(f"\n== iteration {iter} ==") # 1. self-play with the installed champion (include/weights_<size>.c) @@ -409,7 +425,7 @@ def run_step( # 2. generate self-play data (one game per process, in parallel) eps = max(eps_min, eps_start * (eps_decay ** (iter - 1))) - csv = root / "data" / f"selfplay-{size}-v{iter:03d}.csv" + csv = make_csv(iter) print(f"selfplay, eps={eps:.4f}, depth {depth}, {games} games on {procs} procs") out_rows: list[str] = [] for out in fan_out( @@ -429,7 +445,9 @@ def run_step( # 3. fine-tune the champion on TD(lambda) targets -> candidate device = next(champ.parameters()).device - records = SelfPlayRecords.load_selfplay(csv, size, device) + csvs = [make_csv(i) for i in range(max(iter - replay_iters, 1), iter + 1)] + print("Loading data from", csvs) + records = SelfPlayRecords.load_merge(csvs, size, device) cand = copy.deepcopy(champ) corr, mse = cand.fit( @@ -498,6 +516,12 @@ def main() -> None: "--games", type=int, default=2000, help="selfplay games per iteration" ) _ = p.add_argument( + "--replay-iters", + type=int, + default=4, + help="number of previous self-play instances to mix in with current dataset during training", + ) + _ = p.add_argument( "--eps-start", type=float, default=0.15, @@ -507,16 +531,7 @@ def main() -> None: ) _ = p.add_argument("--eps-min", type=float, default=0.05, help="epsilon floor") _ = p.add_argument( - "--eps-decay", type=float, default=0.8, help="per-iteration epsilon decay" - ) - _ = p.add_argument( - "--depth", - type=int, - default=4, - help="search depth for selfplay AND the tournament", - ) - _ = p.add_argument( - "--min-ply", type=int, default=8, help="record positions at this ply or later" + "--eps-decay", type=float, default=0.9, help="per-iteration epsilon decay" ) _ = p.add_argument("--td-lambda", type=float, default=0.6, help="TD(lambda) lambda") _ = p.add_argument( @@ -554,6 +569,16 @@ def main() -> None: default=None, help="parallel game processes (half the CPUs if unset)", ) + _ = p.add_argument( + "--depth", + type=int, + default=4, + help="search depth for selfplay and the tournament", + ) + _ = p.add_argument( + "--min-ply", type=int, default=4, help="record positions at this ply or later" + ) + args = p.parse_args() if args.tournament_games is None: @@ -565,25 +590,26 @@ def main() -> None: champ = setup(size=args.size, root=root) for iter in range(1, args.iters + 1): champ = run_step( - iter=iter, - size=args.size, - eps_min=args.eps_min, + batch=args.batch, champ=champ, - eps_start=args.eps_start, - eps_decay=args.eps_decay, depth=args.depth, - min_ply=args.min_ply, + epochs=args.epochs, + eps_decay=args.eps_decay, + eps_min=args.eps_min, + eps_start=args.eps_start, games=args.games, + iter=iter, + lr=args.lr, + min_ply=args.min_ply, procs=args.procs, + promote_rate=args.promote, + replay_iters=args.replay_iters, + root=root, + size=args.size, td_lambda=args.td_lambda, - val_frac=args.val_frac, - epochs=args.epochs, - batch=args.batch, - lr=args.lr, tournament_eps=args.tournament_eps, tournament_games=args.tournament_games, - promote_rate=args.promote, - root=root, + val_frac=args.val_frac, ) print(f"\n== done, champion in include/weights_{args.size}.c ==") |