move data generation into the epoch loop

Previously we would spend all of the training epochs on making the network a better predictor of the *initial* TD-leaf values that we computed before the start of the first epoch. There's definitely something stale about this, given that the TD-leaf training targets involved the net evaluation at PVs. I'm not overly confident that this makes sense, but as a kind of compromise between fully on-line TD-leaf and the fully stale version we had before, we now recompute the TD-leaf values in each epoch with the fine-tuned net instead of the epoch 0 net.

Perhaps we were just lucky, but this seems to train a strong network under the default setting.

More experimentation is required.

1 files changed, 52 insertions, 52 deletions

diff --git a/src/nn_train.py b/src/nn_train.py
index 7414fab..5332aa8 100755
--- a/src/nn_train.py
+++ b/src/nn_train.py
@@ -115,6 +115,9 @@ def td_lambda_targets(
     return targets
 
 
+SPOOKY: int = 3
+
+
 @dataclass
 class Performance:
     win: int
@@ -125,8 +128,8 @@ class Performance:
         return float(self.win) / float(max(1, self.win + self.loss + self.draw))
 
     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)}"
+        global SPOOKY
+        return f"+{str(self.win).zfill(SPOOKY)}-{str(self.loss).zfill(SPOOKY)}={str(self.draw).zfill(SPOOKY)}"
 
 
 @dataclass(frozen=True)
@@ -288,31 +291,29 @@ class Net(Module):
         batch: int,
         lr: float,
     ) -> tuple[float, float]:
-        # data
-        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))
-        val_idx = perm[:n_val]
-        train_idx = perm[n_val:]
-
-        val_positions = positions[val_idx]
-        val_targets = targets[val_idx]
-        val_parity_signs = parity_signs[val_idx]
-        train_positions = positions[train_idx]
-        train_targets = targets[train_idx]
-        train_parity_signs = parity_signs[train_idx]
-
         # training widgets
         optimizer = torch.optim.AdamW(self.parameters(), lr=lr)
         loss_fn = MSELoss()
 
-        best_corr = -1.0
-        best_state: dict[str, Tensor] | None = None
-        best_mse = float("nan")
-
         pbar = tqdm(range(epochs), desc="train", unit="ep")
         for _ in pbar:
+            # make data
+            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))
+            val_idx = perm[:n_val]
+            train_idx = perm[n_val:]
+
+            val_positions = positions[val_idx]
+            val_targets = targets[val_idx]
+            val_parity_signs = parity_signs[val_idx]
+            train_positions = positions[train_idx]
+            train_targets = targets[train_idx]
+            train_parity_signs = parity_signs[train_idx]
+
             self.train()
             order = torch.randperm(len(train_positions), device=positions.device)
             for start in range(0, len(train_positions), batch):
@@ -335,17 +336,7 @@ class Net(Module):
 
             pbar.set_postfix(val_corr=f"{corr:.4f}", val_mse=f"{mse:.4f}")
 
-            if corr > best_corr:
-                best_corr = corr
-                best_state = {
-                    name: param.clone() for name, param in self.state_dict().items()
-                }
-                best_mse = mse
-
-        if best_state is not None:
-            self.load_state_dict(best_state)
-
-        return best_corr, best_mse
+        return corr, mse
 
     def to_c(
         self,
@@ -507,15 +498,29 @@ def run_step(
     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 = PlayRecords.load_merge(csvs, size)
-    print("candidates have", [f"{r.shape[0]:,}" for r in records.results], "rows")
+    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)
     ]
 
+    print("Tournament results:")
     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}%")
+        print(f"\tcandidate {idx:03d} {perf} {perf.rate() * 100.0:.2f}%")
+
+    # 4. write champ
+    idx, champ, perf = max(results, key=lambda t: t[2].rate())
+    path = root / "include" / f"weights_{size}.c"
+    champ.to_c(
+        size=size,
+        id=None,
+        path=path,
+        note=f"{iter=}, perf={str(perf)}, rate={perf.rate() * 100.0:.3f}",
+    )
+    _ = 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)}")
 
-    # 4. fine-tune on TD(lambda) targets -> candidate
+    # 5. fine-tune on TD-leaf(lambda) targets -> candidate
     for idx, cand in enumerate(candidates):
         corr, mse = cand.fit(
             records=records.index(idx),
@@ -532,19 +537,6 @@ def run_step(
             note=f"id={idx}, v{iter:03d}, val_corr={corr:.4f}, val_mse={mse:.4f}",
         )
 
-    # 5. write champ
-    idx, champ, perf = max(results, key=lambda t: t[2].rate())
-    path = root / "include" / f"weights_{size}.c"
-    champ.to_c(
-        size=size,
-        id=None,
-        path=path,
-        note=f"{iter=}, perf={str(perf)}, rate={perf.rate() * 100.0:.3f}",
-    )
-    _ = 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=} to {str(path)}")
-
 
 # ===================================================================
 # Main
@@ -567,25 +559,25 @@ def main() -> None:
     _ = p.add_argument(
         "--iters",
         type=int,
-        default=20,
+        default=100,
         help="max loop iterations",
     )
     _ = p.add_argument(
         "--games",
         type=int,
-        default=200,
+        default=500,
         help="selfplay games per iteration",
     )
     _ = p.add_argument(
         "--num-candidates",
         type=int,
-        default=5,
+        default=3,
         help="number of distinct candidates to train",
     )
     _ = p.add_argument(
         "--replay-iters",
         type=int,
-        default=4,
+        default=0,
         help="number of previous data sources to retain",
     )
     _ = p.add_argument(
@@ -642,6 +634,14 @@ 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,
+        )
+    )
+
     (root / "data").mkdir(exist_ok=True)
     candidates = setup(num_candidates=args.num_candidates, size=args.size, root=root)