summaryrefslogtreecommitdiff
path: root/annotate_pgns/nnue_model.py
diff options
context:
space:
mode:
authorScott Gasch <[email protected]>2026-10-15 13:11:22 -0700
committerScott Gasch <[email protected]>2026-10-15 13:11:22 -0700
commit993b6eeb2a9ac7b659c473b6883fa53af1ad3a0b (patch)
treec5388824dec15a7691ce9952a6eb4af66e837127 /annotate_pgns/nnue_model.py
parent53b380c658f828c4c0e096be2763e6b012c23abd (diff)
Phase A NNUE: feature-indexing spec, annotation/training pipeline, dump nnuefeatures
Builds out the NNUE Phase A prototype per docs/NNUE.md: - src/nnuefeatures.c: standalone HalfKAv2_hm feature-index reference implementation (new `dump nnuefeatures` command), sourced from Stockfish's half_ka_v2_hm.{h,cpp} (sf_16 tag), including the 0x88->64 square conversion typhoon's own board representation needs. - annotate_pgns/: sharded TWIC/self-play annotation harness (annotate.py), position extraction, quiet-position filtering (flat and game-shaped, plus an upstream --prefilter-quiet option), train/holdout game-level stratified split, WDL-sigmoid K fitting, the Python HalfKAv2_hm reference encoder (nnue_feature_spec.py) and its bit-for-bit cross-checker against the C side (cross_check_features.py), and a first PyTorch training pipeline (nnue_model.py, nnue_train.py, build_training_tensors.py) -- a 256x2-32-32-1-ish net trained locally on CPU against ~1.85M quiet TWIC positions as a pipeline/framework validation pass. - Finish an in-progress InitMagic/InitLMRTable/InitEval -> InitializeBitboardMagic/InitializeLMRTable/InitializeEval rename: chess.h's prototypes hadn't been updated to match the already-renamed definitions/call sites, leaving 4 implicit-function-declaration warnings. - tests/arasan.epd: new curated EPD suite. - docs/: add the project's design docs (previously untracked). Co-Authored-By: Claude Sonnet 5 <[email protected]>
Diffstat (limited to 'annotate_pgns/nnue_model.py')
-rw-r--r--annotate_pgns/nnue_model.py69
1 files changed, 69 insertions, 0 deletions
diff --git a/annotate_pgns/nnue_model.py b/annotate_pgns/nnue_model.py
new file mode 100644
index 0000000..acdc8fd
--- /dev/null
+++ b/annotate_pgns/nnue_model.py
@@ -0,0 +1,69 @@
+#!/usr/bin/env python3
+"""The Phase A prototype net: a small, known HalfKAv2_hm architecture
+per NNUE.md's architecture decisions -- 256x2-32-32-1-ish, a single
+shared feature-transformer weight matrix applied independently to
+each perspective's active-feature set (that sharing, not two
+separate tables, is what makes the two 256-wide "accumulator" halves
+comparable before concatenation), own-perspective-half-first
+concatenation keyed by side-to-move so the net's single scalar
+output is side-to-move-relative by construction (per the "Output
+convention" architecture decision), and a sigmoid-WDL loss against
+the already-fitted K=1.3123 (per "WDL-sigmoid target conversion
+built and run", fit_k_and_label.py) rather than raw centipawn MSE.
+
+Deliberately NOT doing yet (explicitly deferred, see NNUE.md): clipped
+activations / quantization-aware scaling (this is a float32-only
+Phase A prototype, no export-for-C-inference step exists yet), the
+lazy-eval PSQT-term fast path, and the super-lazy material-only
+prefilter. All irrelevant to proving the training pipeline works.
+"""
+import torch
+import torch.nn as nn
+
+from nnue_feature_spec import DIMENSIONS
+
+L1 = 256 # feature-transformer output width (per perspective)
+L2 = 32
+L3 = 32
+
+DEFAULT_K = 1.3123 # pawns; fit_k_and_label.py's Phase-A train-split fit
+
+
+class NnueNet(nn.Module):
+ def __init__(self, dimensions=DIMENSIONS, l1=L1, l2=L2, l3=L3):
+ super().__init__()
+ # Single shared feature transformer: one [dimensions, l1]
+ # weight matrix, applied once per perspective to that
+ # perspective's own active-feature set. mode="sum" matches
+ # Stockfish's own feature-transformer semantics (accumulate
+ # one weight row per active feature).
+ self.ft = nn.EmbeddingBag(dimensions, l1, mode="sum")
+ self.ft_bias = nn.Parameter(torch.zeros(l1))
+
+ self.fc1 = nn.Linear(2 * l1, l2)
+ self.fc2 = nn.Linear(l2, l3)
+ self.fc3 = nn.Linear(l3, 1)
+
+ def forward(self, white_idx, white_off, black_idx, black_off, stm):
+ """stm: float/long tensor of shape [B], 1.0 if white to move
+ else 0.0. Returns raw score in pawn-ish units, shape [B]."""
+ white_acc = self.ft(white_idx, white_off) + self.ft_bias
+ black_acc = self.ft(black_idx, black_off) + self.ft_bias
+
+ white_acc = torch.relu(white_acc)
+ black_acc = torch.relu(black_acc)
+
+ stm = stm.view(-1, 1).float()
+ # own-perspective-first concatenation, per side to move.
+ own = stm * white_acc + (1.0 - stm) * black_acc
+ enemy = stm * black_acc + (1.0 - stm) * white_acc
+ x = torch.cat([own, enemy], dim=1)
+
+ x = torch.relu(self.fc1(x))
+ x = torch.relu(self.fc2(x))
+ x = self.fc3(x)
+ return x.view(-1)
+
+ def win_probability(self, white_idx, white_off, black_idx, black_off, stm, k=DEFAULT_K):
+ raw = self.forward(white_idx, white_off, black_idx, black_off, stm)
+ return torch.sigmoid(raw / k)