| !pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121 |
| """ |
| Chess AI โ AlphaZero-style |
| Backend : PyTorch CUDA (thay thแบฟ TinyGrad) |
| Checkpoint : pickle (giแปฏ nguyรชn) |
| Deps : torch, chess, numpy, tqdm |
| """ |
|
|
| EVAL_CACHE = {} |
| CACHE_MAX = 100_000 |
|
|
| import os, pickle, time, traceback |
| import chess, chess.pgn |
| import numpy as np |
| from tqdm.auto import tqdm |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from torch.optim import Adam |
|
|
| DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| print(f"[Device] {DEVICE} | " |
| f"{torch.cuda.get_device_name(0) if DEVICE.type == 'cuda' else 'CPU'}") |
|
|
| |
| |
| |
| CFG = { |
| "mcts_sims": 16, |
| "buffer_maxlen": 50_000, |
| "min_buffer": 2_000, |
| "batch_size": 512, |
| "lr": 1e-4, |
| "train_steps": 30, |
| "save_every": 1, |
| "pgn_limit": 2000, |
| "res_blocks": 6, |
| "channels": 256, |
| "max_moves": 81, |
| "checkpoint_path": "/content/drive/MyDrive/dataset stockfish/latest.pkl", |
| "pgn_path": "/content/drive/MyDrive/dataset stockfish/million_games.pgn", |
| "log_path": "/content/drive/MyDrive/dataset stockfish/training_log.txt", |
| } |
| CFG["pgn_chunk"] = 500 |
| CFG["pgn_offset_file"] = "/content/drive/MyDrive/dataset stockfish/pgn_offset.pkl" |
| TOTAL_STEPS_TARGET = CFG["train_steps"] |
| os.makedirs(os.path.dirname(CFG["checkpoint_path"]), exist_ok=True) |
|
|
| |
| |
| |
| MCTS_BATCH = 16 |
|
|
| PIECE_IDX = { |
| chess.PAWN: 0, chess.KNIGHT: 1, chess.BISHOP: 2, |
| chess.ROOK: 3, chess.QUEEN: 4, chess.KING: 5, |
| } |
| PROMOTION_OFFSET = 4096 |
| PROMOTION_MAP = {chess.QUEEN: 0, chess.ROOK: 1, |
| chess.BISHOP: 2, chess.KNIGHT: 3} |
| ACTION_SIZE = 4672 |
|
|
| ENCODE_CACHE = {} |
|
|
| def encode_board(board: chess.Board) -> np.ndarray: |
| key = board.fen() |
| if key in ENCODE_CACHE: |
| return ENCODE_CACHE[key] |
| x = np.zeros((12, 8, 8), dtype=np.float32) |
| for sq, piece in board.piece_map().items(): |
| row, col = divmod(sq, 8) |
| offset = 0 if piece.color == chess.WHITE else 6 |
| x[PIECE_IDX[piece.piece_type] + offset][row][col] = 1.0 |
| ENCODE_CACHE[key] = x |
| return x |
|
|
| def move_to_idx(move: chess.Move) -> int: |
| if move.promotion is not None: |
| return PROMOTION_OFFSET + move.to_square * 4 + PROMOTION_MAP[move.promotion] |
| return move.from_square * 64 + move.to_square |
|
|
| |
| |
| |
| class ResBlock(nn.Module): |
| def __init__(self, channels: int): |
| super().__init__() |
| self.c1 = nn.Conv2d(channels, channels, 3, padding=1, bias=False) |
| self.bn1 = nn.BatchNorm2d(channels) |
| self.c2 = nn.Conv2d(channels, channels, 3, padding=1, bias=False) |
| self.bn2 = nn.BatchNorm2d(channels) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| residual = x |
| x = F.relu(self.bn1(self.c1(x))) |
| x = self.bn2(self.c2(x)) |
| return F.relu(x + residual) |
|
|
|
|
| class ChessNet(nn.Module): |
| def __init__(self, channels: int = 256, num_blocks: int = 6): |
| super().__init__() |
| self.stem_conv = nn.Conv2d(12, channels, 3, padding=1, bias=False) |
| self.stem_bn = nn.BatchNorm2d(channels) |
| self.blocks = nn.ModuleList([ResBlock(channels) for _ in range(num_blocks)]) |
|
|
| |
| self.policy_conv = nn.Conv2d(channels, 32, 1, bias=False) |
| self.policy_bn = nn.BatchNorm2d(32) |
| self.policy_fc = nn.Linear(32 * 8 * 8, ACTION_SIZE) |
|
|
| |
| self.value_conv = nn.Conv2d(channels, 8, 1, bias=False) |
| self.value_bn = nn.BatchNorm2d(8) |
| self.value_fc1 = nn.Linear(8 * 8 * 8, 256) |
| self.value_fc2 = nn.Linear(256, 1) |
|
|
| def forward(self, x: torch.Tensor): |
| x = F.relu(self.stem_bn(self.stem_conv(x))) |
| for block in self.blocks: |
| x = block(x) |
|
|
| p = F.relu(self.policy_bn(self.policy_conv(x))) |
| p = self.policy_fc(p.flatten(1)) |
|
|
| v = F.relu(self.value_bn(self.value_conv(x))) |
| v = F.relu(self.value_fc1(v.flatten(1))) |
| v = torch.tanh(self.value_fc2(v)) |
| return p, v |
|
|
|
|
| model = ChessNet(CFG["channels"], CFG["res_blocks"]).to(DEVICE) |
| optimizer = Adam(model.parameters(), lr=CFG["lr"]) |
|
|
| |
| |
| |
| class ReplayBuffer: |
| def __init__(self, maxlen: int): |
| self.maxlen = maxlen |
| self.states = np.zeros((maxlen, 12, 8, 8), dtype=np.float32) |
| self.policies = np.zeros((maxlen, ACTION_SIZE), dtype=np.float32) |
| self.values = np.zeros((maxlen, 1), dtype=np.float32) |
| self.ptr = 0 |
| self.size = 0 |
|
|
| def extend(self, samples): |
| for state, pi, z in samples: |
| self.states[self.ptr] = state |
| self.policies[self.ptr] = pi |
| self.values[self.ptr, 0] = z |
| self.ptr = (self.ptr + 1) % self.maxlen |
| self.size = min(self.size + 1, self.maxlen) |
|
|
| def sample_batch(self, batch_size: int): |
| idx = np.random.randint(0, self.size, batch_size) |
| return self.states[idx], self.policies[idx], self.values[idx] |
|
|
| def __len__(self): |
| return self.size |
|
|
| |
| |
| |
| def save_checkpoint(games_played: int): |
| data = { |
| "games": games_played, |
| "total_steps": total_steps, |
| "weights": {k: v.cpu() for k, v in model.state_dict().items()}, |
| "optimizer": optimizer.state_dict(), |
| } |
| with open(CFG["checkpoint_path"], "wb") as f: |
| pickle.dump(data, f) |
| tqdm.write(f" [CKPT] saved โ games={games_played}, steps={total_steps:,}") |
|
|
|
|
| def load_checkpoint(): |
| path = CFG["checkpoint_path"] |
| if not os.path.exists(path): |
| tqdm.write(" [CKPT] khรดng tรฌm thแบฅy checkpoint, bแบฏt ฤแบงu tแปซ ฤแบงu.") |
| return 0, 0 |
| with open(path, "rb") as f: |
| data = pickle.load(f) |
| model.load_state_dict(data["weights"]) |
| optimizer.load_state_dict(data["optimizer"]) |
| games = data.get("games", 0) |
| steps = data.get("total_steps", 0) |
| tqdm.write(f" [CKPT] loaded โ games={games}, steps={steps:,}") |
| return games, steps |
|
|
| |
| |
| |
| def load_pgn_offset(): |
| if os.path.exists(CFG["pgn_offset_file"]): |
| with open(CFG["pgn_offset_file"], "rb") as f: |
| return pickle.load(f) |
| return 0 |
|
|
|
|
| def load_pgn_chunk(buf: ReplayBuffer, path: str, max_games: int = 500): |
| if not os.path.exists(path): |
| return 0, None |
| offset = load_pgn_offset() |
| count = 0 |
| with open(path, "r") as f: |
| f.seek(offset) |
| while count < max_games: |
| game = chess.pgn.read_game(f) |
| if game is None: |
| break |
| result = game.headers.get("Result", "1/2-1/2") |
| white_z = 1.0 if result == "1-0" else -1.0 if result == "0-1" else 0.0 |
| board = game.board() |
| samples = [] |
| for move in game.mainline_moves(): |
| pi = np.zeros(ACTION_SIZE, dtype=np.float32) |
| pi[move_to_idx(move)] = 1.0 |
| z = white_z if board.turn == chess.WHITE else -white_z |
| samples.append((encode_board(board), pi, np.float32(z))) |
| board.push(move) |
| buf.extend(samples) |
| count += 1 |
| new_offset = f.tell() |
| with open(CFG["pgn_offset_file"], "wb") as f: |
| pickle.dump(new_offset, f) |
| tqdm.write(f"[PGN] loaded chunk={count}, buffer={len(buf):,}") |
| return count, new_offset |
|
|
|
|
| def load_pgn(buf: ReplayBuffer, path: str, limit: int = 2000): |
| if not os.path.exists(path): |
| tqdm.write(f" [PGN] khรดng tรฌm thแบฅy {path}, bแป qua.") |
| return |
| count = 0 |
| with open(path) as f: |
| pbar = tqdm(total=limit, desc="Loading PGN") |
| while count < limit: |
| game = chess.pgn.read_game(f) |
| if game is None: |
| break |
| result = game.headers.get("Result", "1/2-1/2") |
| white_z = 1.0 if result == "1-0" else -1.0 if result == "0-1" else 0.0 |
| board = game.board() |
| samples = [] |
| for move in game.mainline_moves(): |
| pi = np.zeros(ACTION_SIZE, dtype=np.float32) |
| pi[move_to_idx(move)] = 1.0 |
| z = white_z if board.turn == chess.WHITE else -white_z |
| samples.append((encode_board(board), pi, np.float32(z))) |
| board.push(move) |
| buf.extend(samples) |
| count += 1 |
| pbar.update(1) |
| pbar.close() |
| tqdm.write(f" [PGN] loaded {count} games, buffer={len(buf):,}") |
|
|
| |
| |
| |
| def _softmax_masked(logits_np: np.ndarray, mask: np.ndarray) -> np.ndarray: |
| p = logits_np.astype(np.float64) |
| p = np.nan_to_num(p, nan=0.0, posinf=0.0, neginf=0.0) |
| p -= p.max() |
| np.exp(p, out=p) |
| p *= mask |
| total = p.sum() |
| if total > 0: |
| p /= total |
| else: |
| p = mask / mask.sum() |
| return p |
|
|
|
|
| def predict(board: chess.Board): |
| fen = board.fen() |
| cached = EVAL_CACHE.get(fen) |
| if cached is not None: |
| return cached |
|
|
| x = torch.from_numpy(encode_board(board)[None]).to(DEVICE) |
| model.eval() |
| with torch.no_grad(): |
| logits, v = model(x) |
|
|
| p_np = logits[0].cpu().numpy() |
| v_scalar = float(v[0, 0].cpu().item()) |
|
|
| mask = np.zeros(ACTION_SIZE, dtype=np.float64) |
| for m in board.legal_moves: |
| mask[move_to_idx(m)] = 1.0 |
|
|
| p = _softmax_masked(p_np, mask) |
|
|
| result = (p, v_scalar) |
| if len(EVAL_CACHE) >= CACHE_MAX: |
| EVAL_CACHE.clear() |
| EVAL_CACHE[fen] = result |
| return result |
|
|
|
|
| def predict_batch(boards): |
| n = len(boards) |
| states = np.empty((n, 12, 8, 8), dtype=np.float32) |
| for i, b in enumerate(boards): |
| states[i] = encode_board(b) |
|
|
| x = torch.from_numpy(states).to(DEVICE) |
| model.eval() |
| with torch.no_grad(): |
| logits, values = model(x) |
|
|
| logits_np = logits.cpu().numpy().astype(np.float64, copy=False) |
| values_np = values.cpu().numpy() |
|
|
| outputs = [] |
| for i, board in enumerate(boards): |
| mask = np.zeros(ACTION_SIZE, dtype=np.float64) |
| for m in board.legal_moves: |
| mask[move_to_idx(m)] = 1.0 |
| p = _softmax_masked(logits_np[i], mask) |
| outputs.append((p.copy(), float(values_np[i, 0]))) |
| return outputs |
|
|
| |
| |
| |
| class MCTSNode: |
| __slots__ = ("board", "prior", "parent", "children", "visits", "total_value") |
|
|
| def __init__(self, board: chess.Board, prior: float = 0.0, parent=None): |
| self.board = board |
| self.prior = prior |
| self.parent = parent |
| self.children = {} |
| self.visits = 0 |
| self.total_value = 0.0 |
|
|
| @property |
| def q(self) -> float: |
| return self.total_value / self.visits if self.visits > 0 else 0.0 |
|
|
| def uct_score(self) -> float: |
| c_puct = 1.5 |
| parent_visits = self.parent.visits if self.parent else 1 |
| u = c_puct * self.prior * (parent_visits ** 0.5) / (1 + self.visits) |
| return self.q + u |
|
|
|
|
| def terminal_value(board: chess.Board) -> float: |
| outcome = board.outcome() |
| if outcome is None or outcome.winner is None: |
| return 0.0 |
| return -1.0 |
|
|
|
|
| def expand_node(node: MCTSNode, policy: np.ndarray): |
| for move in node.board.legal_moves: |
| if move not in node.children: |
| child_board = node.board.copy(stack=False) |
| child_board.push(move) |
| node.children[move] = MCTSNode( |
| child_board, prior=float(policy[move_to_idx(move)]), parent=node) |
|
|
|
|
| def backup_path(path, value: float): |
| for n in reversed(path): |
| n.visits += 1 |
| n.total_value += value |
| value = -value |
|
|
|
|
| def mcts_search(root: MCTSNode, num_sims: int = 64): |
| remaining = num_sims |
| while remaining > 0: |
| batch_size = min(MCTS_BATCH, remaining) |
| leaves, paths = [], [] |
|
|
| for _ in range(batch_size): |
| node = root |
| path = [node] |
| while node.children and not node.board.is_game_over(): |
| node = max(node.children.values(), key=lambda n: n.uct_score()) |
| path.append(node) |
| leaves.append(node) |
| paths.append(path) |
|
|
| boards_to_eval, eval_indices = [], [] |
| values = [0.0] * batch_size |
| policies = [None] * batch_size |
|
|
| for i, node in enumerate(leaves): |
| if node.board.is_game_over(): |
| values[i] = terminal_value(node.board) |
| else: |
| boards_to_eval.append(node.board) |
| eval_indices.append(i) |
|
|
| if boards_to_eval: |
| outputs = predict_batch(boards_to_eval) |
| for idx, (policy, value) in zip(eval_indices, outputs): |
| policies[idx] = policy |
| values[idx] = value |
|
|
| for i, node in enumerate(leaves): |
| if not node.board.is_game_over(): |
| policy = policies[i] |
| if node == root: |
| legal = list(node.board.legal_moves) |
| if legal: |
| noise = np.random.dirichlet([0.3] * len(legal)) |
| for j, move in enumerate(legal): |
| k = move_to_idx(move) |
| policy[k] = 0.75 * policy[k] + 0.25 * noise[j] |
| expand_node(node, policy) |
| backup_path(paths[i], values[i]) |
|
|
| remaining -= batch_size |
|
|
| |
| |
| |
| def play_one_game(temperature: float = 1.0, game_id: int = 0): |
| board = chess.Board() |
| history = [] |
| move_num = 0 |
|
|
| move_bar = tqdm( |
| total=CFG["max_moves"], desc=f" โ Game #{game_id}", |
| leave=False, dynamic_ncols=True) |
|
|
| try: |
| while not board.is_game_over() and move_num < CFG["max_moves"]: |
| root = MCTSNode(board.copy(stack=False)) |
| mcts_search(root, CFG["mcts_sims"]) |
|
|
| moves_list = list(root.children.keys()) |
| if not moves_list: |
| break |
|
|
| visits = np.array( |
| [root.children[m].visits for m in moves_list], dtype=np.float32) |
| s = visits.sum() |
| if s <= 0 or np.isnan(s): |
| visits[:] = 1.0 |
|
|
| idxs = np.array([move_to_idx(m) for m in moves_list], dtype=np.int32) |
| pi = np.zeros(ACTION_SIZE, dtype=np.float32) |
| pi[idxs] = visits |
| s = pi.sum() |
| if s > 0 and not np.isnan(s): |
| pi /= s |
|
|
| history.append((encode_board(board), pi, board.turn)) |
|
|
| if move_num < 30 and temperature > 0: |
| probs = visits ** (1.0 / temperature) |
| s = probs.sum() |
| if s <= 0 or np.isnan(s): |
| probs = np.ones(len(visits), dtype=np.float32) |
| probs /= probs.sum() |
| move = moves_list[np.random.choice(len(moves_list), p=probs)] |
| else: |
| move = moves_list[int(np.argmax(visits))] |
|
|
| board.push(move) |
| move_num += 1 |
| move_bar.update(1) |
| move_bar.set_postfix( |
| turn="W" if board.turn == chess.BLACK else "B", |
| move=move.uci(), sims=CFG["mcts_sims"]) |
| finally: |
| move_bar.close() |
|
|
| result = board.result() if board.is_game_over() else "1/2-1/2" |
| white_z = 1.0 if result == "1-0" else -1.0 if result == "0-1" else 0.0 |
| samples = [ |
| (state, pi, np.float32(white_z if color == chess.WHITE else -white_z)) |
| for state, pi, color in history |
| ] |
| return samples, result |
|
|
| |
| |
| |
| def train_step(buf: ReplayBuffer): |
| if len(buf) < CFG["min_buffer"]: |
| return None |
|
|
| model.train() |
| train_bar = tqdm( |
| total=CFG["train_steps"], desc=" ๐ง Training", |
| leave=False, dynamic_ncols=True) |
|
|
| loss = loss_policy = loss_value = 0.0 |
|
|
| for _ in range(CFG["train_steps"]): |
| s, p, z = buf.sample_batch(CFG["batch_size"]) |
|
|
| S = torch.from_numpy(np.ascontiguousarray(s)).to(DEVICE) |
| P = torch.from_numpy(np.ascontiguousarray(p)).to(DEVICE) |
| Z = torch.from_numpy(np.ascontiguousarray(z)).to(DEVICE) |
|
|
| logits, v = model(S) |
|
|
| |
| log_prob = F.log_softmax(logits, dim=1) |
| loss_policy = -(P * log_prob).sum(dim=1).mean() |
|
|
| |
| loss_value = F.mse_loss(v, Z) |
|
|
| loss = loss_policy + loss_value |
|
|
| optimizer.zero_grad() |
| loss.backward() |
| optimizer.step() |
|
|
| train_bar.update(1) |
| train_bar.set_postfix(loss=f"{loss.item():.4f}") |
|
|
| train_bar.close() |
|
|
| return { |
| "loss": loss.item(), |
| "loss_policy": loss_policy.item(), |
| "loss_value": loss_value.item(), |
| "steps": CFG["train_steps"], |
| } |
|
|
| |
| |
| |
| if __name__ == "__main__": |
| step_since_pgn = 0 |
|
|
| print("=" * 60) |
| print(" Chess AI โ AlphaZero-style | PyTorch CUDA") |
| print(f" Device : {DEVICE}") |
| print("=" * 60) |
|
|
| buf = ReplayBuffer(CFG["buffer_maxlen"]) |
| games_played, total_steps = load_checkpoint() |
|
|
| if games_played == 0: |
| load_pgn(buf, CFG["pgn_path"], CFG["pgn_limit"]) |
| load_pgn_chunk(buf, CFG["pgn_path"], max_games=500) |
|
|
| print(f"\n Bแบฏt ฤแบงu training tแปซ game #{games_played + 1}") |
| print(f" Buffer: {len(buf):,} | Min train: {CFG['min_buffer']:,}\n") |
|
|
| game_bar = tqdm(initial=games_played, unit="game", |
| desc="๐ฎ Games", dynamic_ncols=True) |
| step_bar = tqdm( |
| total=TOTAL_STEPS_TARGET, initial=total_steps, |
| unit="step", desc="Steps", dynamic_ncols=True, |
| bar_format="{desc}: {n_fmt}/{total_fmt} |{bar}| {elapsed}<{remaining}") |
|
|
| while True: |
| t0 = time.time() |
| step_since_pgn += CFG["train_steps"] |
| if step_since_pgn >= CFG["pgn_chunk"]: |
| load_pgn_chunk(buf, CFG["pgn_path"], CFG["pgn_chunk"]) |
| step_since_pgn = 0 |
|
|
| |
| try: |
| samples, result = play_one_game(temperature=1.0, game_id=games_played + 1) |
| except Exception as e: |
| ts = time.strftime("%Y-%m-%d %H:%M:%S") |
| msg = f"[ERROR] self-play: {e}\n{traceback.format_exc()}" |
| tqdm.write(f"[{ts}] {msg}") |
| try: |
| with open(CFG["log_path"], "a") as f: |
| f.write(f"[{ts}] {msg}\n") |
| except Exception: |
| pass |
| continue |
|
|
| buf.extend(samples) |
| games_played += 1 |
| elapsed_game = time.time() - t0 |
|
|
| |
| t1 = time.time() |
| info = train_step(buf) |
| elapsed_train = time.time() - t1 |
|
|
| if info: |
| step_bar.update(1) |
| total_steps += 1 |
| step_bar.set_postfix( |
| progress=f"{total_steps}/{TOTAL_STEPS_TARGET}", |
| loss=f"{info['loss']:.4f}", |
| p=f"{info['loss_policy']:.4f}", |
| v=f"{info['loss_value']:.4f}", |
| t=f"{elapsed_train:.1f}s", |
| ) |
|
|
| game_bar.update(1) |
| game_bar.set_postfix( |
| result=result, moves=len(samples), |
| buf=f"{len(buf):,}", |
| loss=f"{info['loss']:.4f}" if info else "warmup", |
| t=f"{elapsed_game:.1f}s", |
| ) |
|
|
| |
| ts = time.strftime("%Y-%m-%d %H:%M:%S") |
| line = ( |
| f"game={games_played:5d} | {result} | moves={len(samples):3d} | " |
| f"buf={len(buf):,} | steps={total_steps:,} | " |
| + (f"loss={info['loss']:.4f} " |
| f"(p={info['loss_policy']:.4f} v={info['loss_value']:.4f})" |
| if info else "warmup") |
| + f" | game={elapsed_game:.1f}s train={elapsed_train:.1f}s" |
| ) |
| try: |
| with open(CFG["log_path"], "a") as f: |
| f.write(f"[{ts}] {line}\n") |
| except Exception: |
| pass |
|
|
| if games_played % CFG["save_every"] == 0: |
| save_checkpoint(games_played) |
