1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
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)
|
