| from collections import deque
|
| import random
|
| import torch
|
| import torch
|
| from engine import GameState
|
| from move_finder import find_best_move_shallow
|
| from infer_nnue import gs_to_nnue_features
|
| from nnue_model import NNUE
|
| from tqdm import tqdm
|
| from infer_nnue import NNUEInfer
|
| NNUE_FEATURES = 32
|
| def pad_features(feats):
|
| if len(feats) < NNUE_FEATURES:
|
| return feats + [0] * (NNUE_FEATURES - len(feats))
|
| return feats[:NNUE_FEATURES]
|
|
|
| import pickle
|
|
|
| def load_pgn_dataset(path):
|
| trajectories = []
|
| current_traj = []
|
|
|
| with open(path, "rb") as f:
|
| while True:
|
| try:
|
| chunk = pickle.load(f)
|
| for item in chunk:
|
| current_traj.append(item)
|
|
|
|
|
| if len(current_traj) > 1 and \
|
| current_traj[-1]["stm"] != current_traj[-2]["stm"]:
|
| trajectories.append(current_traj)
|
| current_traj = []
|
|
|
| except EOFError:
|
| break
|
|
|
| if current_traj:
|
| trajectories.append(current_traj)
|
|
|
| return trajectories
|
|
|
|
|
| @torch.no_grad()
|
| @torch.no_grad()
|
| def td_targets_from_traj(model, traj, gamma=0.99):
|
| if len(traj) == 1:
|
| return [0.0]
|
|
|
| feats = [pad_features(x["features"]) for x in traj]
|
| stm = [x["stm"] for x in traj]
|
|
|
| feats = torch.tensor(feats, dtype=torch.long, device="cuda")
|
| stm = torch.tensor(stm, dtype=torch.long, device="cuda")
|
|
|
| values = model(feats, stm).view(-1)
|
|
|
| targets = torch.empty_like(values)
|
|
|
|
|
| targets[:-1] = gamma * (-values[1:])
|
| targets[-1] = values[-1].detach()
|
|
|
|
|
| targets = torch.clamp(targets, -1.0, 1.0)
|
|
|
| return targets.cpu().tolist()
|
|
|
|
|
|
|
| from collections import deque
|
| import random
|
|
|
| class ReplayBuffer:
|
| def __init__(self, capacity=300_000):
|
| self.buf = deque(maxlen=capacity)
|
|
|
| def add(self, f, stm, t):
|
| self.buf.append((f, stm, t))
|
|
|
| def sample(self, n):
|
| return random.sample(self.buf, n)
|
|
|
| def __len__(self):
|
| return len(self.buf)
|
|
|
|
|
| def train_from_replay(model, optimizer, replay, batch_size):
|
| if len(replay) < batch_size:
|
| return
|
|
|
| batch = replay.sample(batch_size)
|
| feats, stm, targets = zip(*batch)
|
|
|
| feats = torch.tensor(feats, dtype=torch.long, device="cuda")
|
| stm = torch.tensor(stm, dtype=torch.long, device="cuda")
|
| targ = torch.tensor(targets, dtype=torch.float, device="cuda")
|
|
|
| preds = model(feats, stm).view(-1)
|
|
|
| loss = torch.nn.functional.smooth_l1_loss(preds, targ)
|
|
|
| optimizer.zero_grad(set_to_none=True)
|
| loss.backward()
|
| optimizer.step()
|
|
|
|
|
| from tqdm import tqdm
|
| device = "cuda"
|
| model = NNUE().to(device)
|
| model.load_state_dict(torch.load("nnue_model.pt", weights_only=True))
|
| optimizer = torch.optim.Adam(model.parameters(), lr=5e-5)
|
|
|
| replay = ReplayBuffer()
|
| trajectories = load_pgn_dataset("nnue_dataset.pkl")
|
|
|
| for epoch in range(3):
|
| print(f"Epoch {epoch}")
|
|
|
| for traj in tqdm(trajectories):
|
| if len(traj) < 2:
|
| continue
|
|
|
| targets = td_targets_from_traj(model, traj)
|
|
|
| for x, t in zip(traj, targets):
|
| replay.add(
|
| pad_features(x["features"]),
|
| x["stm"],
|
| t
|
| )
|
|
|
| for _ in range(3):
|
| train_from_replay(model, optimizer, replay, batch_size=512)
|
|
|
| torch.save(model.state_dict(), "nnue_model_td.pt")
|
|
|
