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
|
#
# 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/>.
import math
from dataclasses import dataclass
from itertools import product
from pathlib import Path
import torch
from torch.nn import MSELoss
from tqdm import tqdm
from ctak_train.data import DataSplit, Pool
from ctak_train.net import NET_DIMS, Net
from ctak_train.selfplay import cand_copy_for
@dataclass(frozen=True)
class DataGenConfig:
size: int
depth: int
min_ply: int
games: int
procs: int
replay_iters: int
root: Path
train_only_self_play: bool
@dataclass(frozen=True)
class TrainingGridSearchConfig:
smooth_factor: list[float]
epochs: list[int]
batch: list[int]
lr: list[float]
val_frac: list[float]
def make_configs(self, data_gen_config: DataGenConfig) -> list["Candidate"]:
inp_dim, hidden = NET_DIMS[data_gen_config.size]
flds = self.__dataclass_fields__.keys()
almost_cands = [
dict(zip(flds, vals)) for vals in product(*(getattr(self, k) for k in flds))
]
for ident, almost_cand in enumerate(almost_cands):
net = Net(inp_dim, hidden)
net.to_c(
size=data_gen_config.size,
ident=ident,
path=data_gen_config.root
/ "include"
/ f"cand_{ident}_weights_{data_gen_config.size}.c",
note=f"{ident=}, v000, random init",
)
cand_copy_for(
6 if data_gen_config.size == 5 else 5, ident, data_gen_config.root
)
almost_cand["net"] = net
almost_cand["train_only_self_play"] = data_gen_config.train_only_self_play
return [Candidate(**d) for d in almost_cands]
@dataclass
class Candidate:
train_only_self_play: bool
smooth_factor: float
epochs: int
batch: int
lr: float
val_frac: float
net: "Net"
def __str__(self) -> str:
return f"(λ={self.smooth_factor}, E={self.epochs}, BS={self.batch}, η={self.lr:.2e}, V%={self.val_frac * 100.0:.1f})"
def training_pool(self, pool: Pool) -> Pool:
if self.train_only_self_play:
return pool.select_chains(pool.chain_self_play)
return pool
def fit(self, pool: Pool) -> tuple[float, float]:
# The train/val split is fixed but targets are refreshed
split = DataSplit.build(self.training_pool(pool), self.val_frac)
if split.n == 0:
return math.nan, math.nan
optimizer = torch.optim.AdamW(self.net.parameters(), lr=self.lr)
loss_fn = MSELoss()
pbar = tqdm(range(self.epochs), desc="train", unit="ep")
for _ in pbar:
split.refresh_targets(self.net, self.smooth_factor)
self.net.train()
for head, cells, parity, target in split.batches(self.batch):
optimizer.zero_grad()
predictions = self.net(head, cells) * parity
loss = loss_fn(predictions, target)
loss.backward()
optimizer.step()
with torch.no_grad():
self.net.eval()
val = split.val_idx
predictions = self.net(split.pool.head[val], split.pool.cells[val])
predictions = predictions * split.pool.parity[val]
targets = split.targets[val]
mse = loss_fn(predictions, targets).item()
corr = torch.corrcoef(torch.stack([predictions, targets]))[0, 1].item()
pbar.set_postfix(val_corr=f"{corr:.4f}", val_mse=f"{mse:.4f}")
return corr, mse
|
