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

1 files changed, 129 insertions, 0 deletions

diff --git a/train/src/ctak_train/net.py b/train/src/ctak_train/net.py
new file mode 100644
index 0000000..cc8bf5c
--- /dev/null
+++ b/train/src/ctak_train/net.py
@@ -0,0 +1,129 @@
+#
+# 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/>.
+
+from pathlib import Path
+
+import numpy as np
+import torch
+from torch import Tensor
+from torch.nn import Embedding, Linear, Module
+
+# (input dim, hidden dim) per board size; must match include/weights_<size>.h
+NET_DIMS: dict[int, tuple[int, int]] = {
+    5: (27, 64),  # FIVE_INP_NUM, FIVE_DENSE_NUM
+    6: (38, 64),  # SIX_INP_NUM, SIX_DENSE_NUM
+}
+
+
+def _fmt_flat(a: np.ndarray, per_line: int = 8) -> str:
+    a = np.ravel(np.asarray(a, dtype=np.float32))
+    lines: list[str] = []
+    for i in range(0, len(a), per_line):
+        chunk = ", ".join(f"{v:.9g}" for v in a[i : i + per_line])
+        if i + per_line < len(a):
+            chunk += ","
+        lines.append("    " + chunk)
+    return "\n".join(lines)
+
+
+def _fmt_matrix(mat: np.ndarray, per_line: int = 8) -> str:
+    mat = np.asarray(mat, dtype=np.float32)
+    rows: list[str] = []
+    for r in mat:
+        chunks = [
+            ", ".join(f"{v:.9g}" for v in r[i : i + per_line])
+            for i in range(0, len(r), per_line)
+        ]
+        rows.append("    {" + ",\n     ".join(chunks) + "},")
+    return "\n".join(rows)
+
+
+def fast_tanh(x: Tensor) -> Tensor:
+    # Must match fast_tanh in include/nn.c: degree-11 fit of tanh on [-3, 3],
+    # saturating to +-1 outside.
+    x2 = x * x
+    poly = x * (
+        0.997136032
+        + x2
+        * (
+            -0.305215124
+            + x2
+            * (
+                0.0832466371
+                + x2 * (-0.0139947924 + x2 * (0.00123106314 + x2 * -4.29748644e-05))
+            )
+        )
+    )
+    return torch.where(x > 3.0, 1.0, torch.where(x < -3.0, -1.0, poly))
+
+
+def fast_silu(x: Tensor) -> Tensor:
+    # Must match fast_silu in include/nn.c: [4,2] Pade approximant with cutoff.
+    # The cutoff below is evaluated at x = 0 (instead of x) because torch
+    # evaluates both branches of where: 1 + 0.1 * x is 0 for x = -10.
+    xs = torch.where(x < -4.6, 0.0, x)
+    poly = xs * (0.5 + xs * (0.25 + xs * (0.05 + 0.004167 * xs))) / (1 + 0.1 * xs)
+    return torch.where(x < -4.6, 0.0, torch.where(x > 4.6, x, poly))
+
+
+class Net(Module):
+    def __init__(self, inp_dim: int, hidden: int) -> None:
+        super().__init__()
+        # learned cell-state embedding, see nn.c
+        self.embedding = Embedding(7, 1)
+        self.hidden = Linear(inp_dim, hidden, bias=True)
+        self.output = Linear(hidden, 1, bias=True)
+
+    def forward(self, head: Tensor, cells: Tensor) -> Tensor:
+        # head: [n, 2] fixed head inputs; cells: [n, n_cells] raw cell states
+        # (0..6) - the only place they become anything other than integers
+        inp = torch.cat([head, self.embedding(cells.long()).flatten(1)], dim=1)
+        hidden = fast_silu(self.hidden(inp))
+        return fast_tanh(self.output(hidden)).squeeze(-1)
+
+    def to_c(
+        self,
+        *,
+        size: int,
+        ident: int | None,
+        path: Path,
+        note: str,
+    ) -> None:
+        upper = "FIVE" if size == 5 else "SIX"
+        prefix = (f"cand_{ident}_" if ident is not None else "") + upper
+        guard = f"WEIGHTS_{upper}" + (f"_{ident}" if ident is not None else "")
+
+        E = self.embedding.weight.detach().cpu().numpy()
+        W = self.hidden.weight.detach().cpu().numpy()
+        B = self.hidden.bias.detach().cpu().numpy()
+        LW = self.output.weight.detach().cpu().numpy().ravel()
+        LB = float(self.output.bias.detach().cpu().item())
+
+        with open(path, "w") as f:
+            lines = [
+                "/* generated by nn_train.py, do not edit by hand */\n",
+                f"/* {note} */\n",
+                f"#ifndef {guard}\n#define {guard}\n",
+                f'#include "weights_{size}.h"\n\n',
+                f"const float {prefix}_embedding[{len(E)}] =\n{{\n{_fmt_flat(E, per_line=7)}\n}};\n\n",
+                f"const float {prefix}_dense_weights[{upper}_INP_NUM][{upper}_DENSE_NUM] NN_ALIGNED =\n{{\n{_fmt_matrix(W.T)}\n}};\n\n",
+                f"const float {prefix}_dense_biases[{upper}_DENSE_NUM] =\n{{\n{_fmt_flat(B)}\n}};\n\n",
+                f"const float {prefix}_logit_weights[{upper}_DENSE_NUM] NN_ALIGNED =\n{{\n{_fmt_flat(LW)}\n}};\n\n",
+                f"const float {prefix}_logit_bias = {LB:.9g};\n",
+                "#endif",
+            ]
+            for line in lines:
+                _ = f.write(line)