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
118
119
120
121
122
123
124
125
126
127
128
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)
|
