diff options
| -rwxr-xr-x | src/nn_train.py | 99 |
1 files changed, 59 insertions, 40 deletions
diff --git a/src/nn_train.py b/src/nn_train.py index a97c964..4de0a5a 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 itertools import product from math import log, ceil import argparse @@ -108,13 +109,32 @@ class DataGenConfig: @dataclass(frozen=True) -class TrainingConfig: +class TrainingGridSearchConfig: + smooth_factor: list[float] + epochs: list[int] + batch: list[int] + lr: list[float] + val_frac: list[float] + + def make_configs(self) -> list[CandidateTrainingConfig]: + flds = self.__dataclass_fields__.keys() + return [ + CandidateTrainingConfig(**dict(zip(flds, vals))) + for vals in product(*(getattr(self, k) for k in flds)) + ] + + +@dataclass(frozen=True) +class CandidateTrainingConfig: smooth_factor: float epochs: int batch: int lr: float val_frac: float + def __str__(self) -> str: + return f"(λ={self.smooth_factor}, E={self.epochs}, BS={self.batch}, η={self.lr}, V%={self.val_frac * 100.0:.1f})" + # =================================================================== # Self-play records @@ -202,7 +222,9 @@ class CandidateRecords: result.draw += 1 return result - def make_training_data(self, model: Module, training: TrainingConfig) -> DataSplit: + 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 @@ -397,18 +419,18 @@ class Net(Module): return tanh(self.output(hidden)).squeeze(-1) def fit( - self, *, records: CandidateRecords, training: TrainingConfig + self, *, records: CandidateRecords, config: CandidateTrainingConfig ) -> tuple[float, float]: # training widgets - optimizer = torch.optim.AdamW(self.parameters(), lr=training.lr) + optimizer = torch.optim.AdamW(self.parameters(), lr=config.lr) loss_fn = MSELoss() - pbar = tqdm(range(training.epochs), desc="train", unit="ep") + pbar = tqdm(range(config.epochs), desc="train", unit="ep") for _ in pbar: - split = records.make_training_data(self, training) + split = records.make_training_data(self, config) self.train() - for positions, targets, parity_signs in split.batches(training.batch): + for positions, targets, parity_signs in split.batches(config.batch): optimizer.zero_grad() predictions = self(positions) * parity_signs loss = loss_fn(predictions, targets) @@ -535,7 +557,7 @@ def run_step( data_gen: DataGenConfig, iter: int, root: Path, - training: TrainingConfig, + training_configs: list[CandidateTrainingConfig], ): def data_path(i: int) -> Path: return root / "data" / f"selfplay-{data_gen.size}-v{i:03d}.json" @@ -595,7 +617,9 @@ def run_step( 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}%") + print( + f"\tcandidate {idx:03d} {perf} {perf.rate() * 100.0:.2f}% {training_configs[idx]}" + ) # 4. write champ idx, champ, perf = max(results, key=lambda t: t[2].rate()) @@ -612,7 +636,7 @@ def run_step( # 5. fine-tune on the hybrid targets -> candidate for idx, cand in enumerate(candidates): - corr, mse = cand.fit(records=records.index(idx), training=training) + corr, mse = cand.fit(records=records.index(idx), config=training_configs[idx]) cand.to_c( size=data_gen.size, id=idx, @@ -648,50 +672,49 @@ def main() -> None: _ = p.add_argument( "--games", type=int, - default=500, - help="selfplay games per iteration", - ) - _ = p.add_argument( - "--num-candidates", - type=int, - default=3, - help="number of distinct candidates to train", + default=200, + help="games per iteration", ) _ = p.add_argument( "--replay-iters", type=int, default=10, - help="number of previous data sources to retain", + help="grid points for number of previous data sources to retain", ) _ = p.add_argument( "--smooth-factor", type=float, - default=0.7, - help="The EMA smoothing factor", + nargs="+", + default=[0.3, 0.5, 0.7, 0.9], + help="grid points for the EMA smoothing factor", ) _ = p.add_argument( "--epochs", type=int, - default=5, - help="candidate training epochs", + nargs="+", + default=[5], + help="grid points for candidate training epochs", ) _ = p.add_argument( "--lr", type=float, - default=1e-4, - help="learning rate", + nargs="+", + default=[1e-4], + help="grid points for learning rate", ) _ = p.add_argument( "--batch", type=int, - default=128, - help="training batch size", + nargs="+", + default=[128], + help="grid points for training batch size", ) _ = p.add_argument( "--val-frac", type=float, - default=0.05, - help="fraction of rows held out for validation", + nargs="+", + default=[0.05], + help="grid points for fraction of rows held out for validation", ) _ = p.add_argument( "--procs", @@ -717,14 +740,6 @@ def main() -> None: if args.procs is None: args.procs = max(1, (os.cpu_count() or 1)) - global SPOOKY - SPOOKY = ceil( - log( - max(1, args.games * args.num_candidates * (args.num_candidates + 1) // 2), - 10, - ) - ) - data_gen = DataGenConfig( size=args.size, depth=args.depth, @@ -733,7 +748,7 @@ def main() -> None: procs=args.procs, replay_iters=args.replay_iters, ) - training = TrainingConfig( + tgsc = TrainingGridSearchConfig( smooth_factor=args.smooth_factor, epochs=args.epochs, batch=args.batch, @@ -741,8 +756,12 @@ def main() -> None: val_frac=args.val_frac, ) + training_configs = tgsc.make_configs() + (root / "data").mkdir(exist_ok=True) - candidates = setup(num_candidates=args.num_candidates, data_gen=data_gen, root=root) + 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) print(f"Number of trainable parameters: {num_params:,}") @@ -753,7 +772,7 @@ def main() -> None: data_gen=data_gen, iter=iter, root=root, - training=training, + training_configs=training_configs, ) print(f"\n== done, champion in include/weights_{args.size}.c ==") |