factor out nn_train.py into a separate python project to make it more manageable, add plotting and offline strength fit

1 files changed, 125 insertions, 0 deletions

diff --git a/train/src/ctak_train/candidate.py b/train/src/ctak_train/candidate.py
new file mode 100644
index 0000000..aa2a6a6
--- /dev/null
+++ b/train/src/ctak_train/candidate.py
@@ -0,0 +1,125 @@
+#
+# 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 math
+from dataclasses import dataclass
+from itertools import product
+from pathlib import Path
+
+import torch
+from torch.nn import MSELoss
+from tqdm import tqdm
+
+from ctak_train.data import DataSplit, Pool
+from ctak_train.net import NET_DIMS, Net
+from ctak_train.selfplay import cand_copy_for
+
+
+@dataclass(frozen=True)
+class DataGenConfig:
+    size: int
+    depth: int
+    min_ply: int
+    games: int
+    procs: int
+    replay_iters: int
+    root: Path
+    train_only_self_play: bool
+
+
+@dataclass(frozen=True)
+class TrainingGridSearchConfig:
+    smooth_factor: list[float]
+    epochs: list[int]
+    batch: list[int]
+    lr: list[float]
+    val_frac: list[float]
+
+    def make_configs(self, data_gen_config: DataGenConfig) -> list["Candidate"]:
+        inp_dim, hidden = NET_DIMS[data_gen_config.size]
+        flds = self.__dataclass_fields__.keys()
+        almost_cands = [
+            dict(zip(flds, vals)) for vals in product(*(getattr(self, k) for k in flds))
+        ]
+        for ident, almost_cand in enumerate(almost_cands):
+            net = Net(inp_dim, hidden)
+            net.to_c(
+                size=data_gen_config.size,
+                ident=ident,
+                path=data_gen_config.root
+                / "include"
+                / f"cand_{ident}_weights_{data_gen_config.size}.c",
+                note=f"{ident=}, v000, random init",
+            )
+            cand_copy_for(
+                6 if data_gen_config.size == 5 else 5, ident, data_gen_config.root
+            )
+            almost_cand["net"] = net
+            almost_cand["train_only_self_play"] = data_gen_config.train_only_self_play
+        return [Candidate(**d) for d in almost_cands]
+
+
+@dataclass
+class Candidate:
+    train_only_self_play: bool
+    smooth_factor: float
+    epochs: int
+    batch: int
+    lr: float
+    val_frac: float
+    net: "Net"
+
+    def __str__(self) -> str:
+        return f"(λ={self.smooth_factor}, E={self.epochs}, BS={self.batch}, η={self.lr:.2e}, V%={self.val_frac * 100.0:.1f})"
+
+    def training_pool(self, pool: Pool) -> Pool:
+        if self.train_only_self_play:
+            return pool.select_chains(pool.chain_self_play)
+        return pool
+
+    def fit(self, pool: Pool) -> tuple[float, float]:
+        # The train/val split is fixed but targets are refreshed
+        split = DataSplit.build(self.training_pool(pool), self.val_frac)
+        if split.n == 0:
+            return math.nan, math.nan
+
+        optimizer = torch.optim.AdamW(self.net.parameters(), lr=self.lr)
+        loss_fn = MSELoss()
+
+        pbar = tqdm(range(self.epochs), desc="train", unit="ep")
+        for _ in pbar:
+            split.refresh_targets(self.net, self.smooth_factor)
+
+            self.net.train()
+            for head, cells, parity, target in split.batches(self.batch):
+                optimizer.zero_grad()
+                predictions = self.net(head, cells) * parity
+                loss = loss_fn(predictions, target)
+                loss.backward()
+                optimizer.step()
+
+            with torch.no_grad():
+                self.net.eval()
+                val = split.val_idx
+                predictions = self.net(split.pool.head[val], split.pool.cells[val])
+                predictions = predictions * split.pool.parity[val]
+                targets = split.targets[val]
+                mse = loss_fn(predictions, targets).item()
+                corr = torch.corrcoef(torch.stack([predictions, targets]))[0, 1].item()
+
+            pbar.set_postfix(val_corr=f"{corr:.4f}", val_mse=f"{mse:.4f}")
+
+        return corr, mse