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
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
|
#!/usr/bin/env python3
"""Precompute HalfKAv2_hm feature tensors from a game-shaped labeled
file (quiet_filter_games.py output) into a cached, CSR-style on-disk
format that nnue_train.py can mmap and sample from directly --
avoids re-parsing FEN / recomputing SEE / recomputing feature indices
on every training epoch (the expensive part, done once here).
Each position's active feature indices (per nnue_feature_spec.py,
already cross-checked bit-for-bit against the C reference -- see
cross_check_features.py) are stored CSR-style (one flat concatenated
index array + a row-offset array per perspective), exactly the input
shape torch.nn.EmbeddingBag expects -- no fixed-width padding needed
since bag sizes vary (3-32 active features per position).
Skips exact rows (no p_search/r_outcome label -- see NNUE.md: a
mate/EGTB/draw-claim score is a game-theoretic fact, not something to
fit a sigmoid-WDL target against).
Usage:
python3 build_training_tensors.py \
--in runs/sd10_twic_2M/games.train.labeled.quiet.jsonl \
--outdir runs/sd10_twic_2M/tensors_train \
--lambda 1.0
"""
import argparse
import array
import json
import sys
from pathlib import Path
import numpy as np
from nnue_feature_spec import active_feature_indices, chess
def iter_positions(path):
with open(path) as f:
for line in f:
game = json.loads(line)
for pos in game["positions"]:
yield pos
def main():
ap = argparse.ArgumentParser(description=__doc__)
ap.add_argument("--in", dest="infile", required=True)
ap.add_argument("--outdir", required=True)
ap.add_argument("--lambda", dest="lam", type=float, default=1.0,
help="target = lambda*p_search + (1-lambda)*r_outcome "
"(Phase A default 1.0, pure p_search -- see NNUE.md)")
ap.add_argument("--limit", type=int, default=None,
help="stop after this many usable positions (smoke testing)")
args = ap.parse_args()
outdir = Path(args.outdir)
outdir.mkdir(parents=True, exist_ok=True)
# array.array('i') is a packed C int buffer -- avoids the ~28
# bytes/element overhead of a Python list of boxed ints, which
# matters once nnz reaches tens of millions.
white_idx = array.array("i")
black_idx = array.array("i")
white_off = [0]
black_off = [0]
stm = [] # 1 = white to move, 0 = black
target = []
n_total = 0
n_used = 0
n_skipped = 0
for pos in iter_positions(args.infile):
n_total += 1
if args.limit is not None and n_used >= args.limit:
break
if pos.get("exact") or "p_search" not in pos:
n_skipped += 1
continue
board = chess.Board(pos["fen"])
w = active_feature_indices(board, chess.WHITE)
b = active_feature_indices(board, chess.BLACK)
white_idx.extend(w)
white_off.append(white_off[-1] + len(w))
black_idx.extend(b)
black_off.append(black_off[-1] + len(b))
stm.append(1 if board.turn == chess.WHITE else 0)
p_search = pos["p_search"]
r_outcome = pos.get("r_outcome", p_search)
target.append(args.lam * p_search + (1.0 - args.lam) * r_outcome)
n_used += 1
if n_used % 200000 == 0:
print(f" ...{n_used} positions processed", file=sys.stderr)
print(f"Read {n_total} rows ({args.infile})")
print(f" {n_skipped} skipped (exact / no p_search label)")
print(f" {n_used} usable positions")
np.save(outdir / "white_idx.npy", np.frombuffer(white_idx, dtype=np.int32))
np.save(outdir / "white_off.npy", np.array(white_off, dtype=np.int64))
np.save(outdir / "black_idx.npy", np.frombuffer(black_idx, dtype=np.int32))
np.save(outdir / "black_off.npy", np.array(black_off, dtype=np.int64))
np.save(outdir / "stm.npy", np.array(stm, dtype=np.uint8))
np.save(outdir / "target.npy", np.array(target, dtype=np.float32))
nnz_white = len(white_idx)
nnz_black = len(black_idx)
print(f"Wrote tensors to {outdir}: {n_used} positions, "
f"{nnz_white} white / {nnz_black} black nonzero feature entries "
f"({nnz_white / max(n_used, 1):.1f} avg features/position)")
if __name__ == "__main__":
main()
|
