Create train.py

train.py

@@ -0,0 +1,555 @@
+def train():
+    """
+    Chess AI (AlphaZero-style) — tối ưu cho Google Colab
+    =====================================================
+    Cách dùng trên Colab:
+      !pip install chess tinygrad
+      from google.colab import drive; drive.mount('/content/drive')
+      !python chess_ai_colab.py
+    """
+    import os
+    import time
+
+    # ==================== ENV ====================
+    os.environ["DEV"] = "CUDA"
+
+    import chess
+    import chess.pgn
+    import numpy as np
+    import pickle
+    from tinygrad import Tensor, nn
+    from tinygrad.nn.state import get_state_dict, load_state_dict, get_parameters
+    from tinygrad.nn.optim import Adam
+
+    def log(msg):
+        print(msg, flush=True)
+
+    # =============================================================
+    # CONFIG
+    # =============================================================
+    CFG = {
+        "mcts_sims":       16,
+        "buffer_maxlen":   52_000,
+        "min_buffer":      1,
+        "batch_size":      512,
+        "lr":              1e-4,
+        "train_steps":     30,
+        "save_every":      10,
+        "pgn_limit":       2000,
+        "res_blocks":      6,
+        "max_moves":       81,
+        "checkpoint_path": "dataset stockfish/latest.pkl",
+        "pgn_path":        "dataset stockfish/million_games.pgn",
+        "log_path":        "dataset stockfish/training_log.txt",
+    }
+    CFG["pgn_chunk"]      = 500
+    CFG["pgn_offset_file"] = "dataset stockfish/pgn_offset.pkl"
+    CFG["channels"]       = 256
+    TOTAL_STEPS_TARGET    = CFG["train_steps"]
+    os.makedirs(os.path.dirname(CFG["checkpoint_path"]), exist_ok=True)
+
+    # =============================================================
+    # HUGGING FACE CONFIG
+    # =============================================================
+    HF_REPO  = "bobbydeveloper2014/bybychessengine"
+    HF_TOKEN = "yourtoken"
+    HF_SYNC_FILES = {
+        "latest.pkl":       CFG["checkpoint_path"],
+        "pgn_offset.pkl":   CFG["pgn_offset_file"],
+        "training_log.txt": CFG["log_path"],
+    }
+
+    hf_api = None
+    try:
+        from huggingface_hub import HfApi, create_repo
+        hf_api = HfApi(token=HF_TOKEN)
+        create_repo(HF_REPO, token=HF_TOKEN, exist_ok=True, repo_type="model")
+        log(f"  [HF] sẵn sàng sync với {HF_REPO}")
+    except Exception as e:
+        log(f"  [HF] lỗi khởi tạo (bỏ qua): {e}")
+
+    def push_to_hf(games_played, total_steps):
+        if hf_api is None:
+            return
+        for remote_name, local_path in HF_SYNC_FILES.items():
+            if not os.path.exists(local_path):
+                continue
+            try:
+                hf_api.upload_file(
+                    path_or_fileobj=local_path,
+                    path_in_repo=remote_name,
+                    repo_id=HF_REPO,
+                    repo_type="model",
+                    commit_message=f"checkpoint @ game={games_played} steps={total_steps}",
+                )
+            except Exception as e:
+                log(f"  [HF][ERROR] upload {remote_name}: {e}")
+        log(f"  [HF] uploaded -> games={games_played}, steps={total_steps:,}")
+
+    # =============================================================
+    # ENCODING
+    # =============================================================
+    MCTS_BATCH       = 16
+    ACTION_SIZE      = 4672
+    PROMOTION_OFFSET = 4096
+    PROMOTION_MAP    = {chess.QUEEN: 0, chess.ROOK: 1, chess.BISHOP: 2, chess.KNIGHT: 3}
+    PIECE_IDX        = {chess.PAWN: 0, chess.KNIGHT: 1, chess.BISHOP: 2,
+                        chess.ROOK: 3, chess.QUEEN: 4, chess.KING: 5}
+    ENCODE_CACHE     = {}
+    EVAL_CACHE       = {}
+    CACHE_MAX        = 100_000
+
+    def encode_board(board):
+        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():
+            r, c = divmod(sq, 8)
+            offset = 0 if piece.color == chess.WHITE else 6
+            x[PIECE_IDX[piece.piece_type] + offset][r][c] = 1.0
+        ENCODE_CACHE[key] = x
+        return x
+
+    def move_to_idx(move):
+        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
+
+    # =============================================================
+    # MODEL
+    # =============================================================
+    class ResBlock:
+        def __init__(self, channels):
+            self.c1  = nn.Conv2d(channels, channels, 3, padding=1, bias=False)
+            self.bn1 = nn.BatchNorm(channels)
+            self.c2  = nn.Conv2d(channels, channels, 3, padding=1, bias=False)
+            self.bn2 = nn.BatchNorm(channels)
+        def __call__(self, x):
+            r = x
+            x = self.bn1(self.c1(x)).relu()
+            x = self.bn2(self.c2(x))
+            return (x + r).relu()
+
+    class ChessNet:
+        def __init__(self, channels=256, num_blocks=6):
+            self.stem_conv   = nn.Conv2d(12, channels, 3, padding=1, bias=False)
+            self.stem_bn     = nn.BatchNorm(channels)
+            self.blocks      = [ResBlock(channels) for _ in range(num_blocks)]
+            self.policy_conv = nn.Conv2d(channels, 32, 1, bias=False)
+            self.policy_bn   = nn.BatchNorm(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.BatchNorm(8)
+            self.value_fc1   = nn.Linear(8 * 8 * 8, 256)
+            self.value_fc2   = nn.Linear(256, 1)
+        def __call__(self, x):
+            x = self.stem_bn(self.stem_conv(x)).relu()
+            for b in self.blocks: x = b(x)
+            p = self.policy_bn(self.policy_conv(x)).relu()
+            p = self.policy_fc(p.flatten(1))
+            v = self.value_bn(self.value_conv(x)).relu()
+            v = self.value_fc1(v.flatten(1)).relu()
+            v = self.value_fc2(v).tanh()
+            return p, v
+
+    model     = ChessNet(CFG["channels"], CFG["res_blocks"])
+    optimizer = Adam(get_parameters(model), lr=CFG["lr"])
+
+    # =============================================================
+    # CHECKPOINT — weights luôn lưu numpy, tương thích Metal + CUDA
+    # =============================================================
+    def _tensor_to_numpy(v):
+        """Convert bất kỳ kiểu weight nào → numpy float32."""
+        if isinstance(v, np.ndarray):
+            return v.astype(np.float32)
+        if isinstance(v, Tensor):
+            try:
+                return v.numpy().astype(np.float32)
+            except Exception:
+                pass
+        if hasattr(v, 'numpy'):
+            try:
+                return np.array(v.numpy(), dtype=np.float32)
+            except Exception:
+                pass
+        if hasattr(v, 'data'):
+            try:
+                return np.array(v.data, dtype=np.float32)
+            except Exception:
+                pass
+        try:
+            return np.array(v, dtype=np.float32)
+        except Exception:
+            return None
+
+    def save_checkpoint(games_played):
+        # Luôn save numpy để macOS Metal load được
+        numpy_weights = {}
+        for k, v in get_state_dict(model).items():
+            arr = _tensor_to_numpy(v)
+            if arr is not None:
+                numpy_weights[k] = arr
+
+        data = {
+            "games":       games_played,
+            "total_steps": total_steps,
+            "weights":     numpy_weights,
+        }
+        with open(CFG["checkpoint_path"], "wb") as f:
+            pickle.dump(data, f)
+        log(f"  [CKPT] saved → games={games_played}, steps={total_steps:,}")
+
+    def load_checkpoint():
+        path = CFG["checkpoint_path"]
+        if not os.path.exists(path):
+            log("  [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)
+
+        # Convert tất cả về Tensor để load vào model
+        clean = {}
+        nan_keys = []
+        for k, v in data["weights"].items():
+            arr = _tensor_to_numpy(v)
+            if arr is not None:
+                if not np.isfinite(arr).all():
+                    nan_keys.append(k)
+                clean[k] = Tensor(arr)
+        load_state_dict(model, clean)
+
+        if nan_keys:
+            log(f"  [CKPT][CẢNH BÁO] checkpoint chứa NaN/Inf ở {len(nan_keys)} tensor: {nan_keys[:5]}{'...' if len(nan_keys) > 5 else ''}")
+            log("  [CKPT] -> checkpoint có thể đã hỏng. Kiểm tra lại trước khi train tiếp.")
+
+        games = data.get("games", 0)
+        steps = data.get("total_steps", 0)
+        log(f"  [CKPT] loaded → games={games}, steps={steps:,}")
+        return games, steps
+
+    # =============================================================
+    # REPLAY BUFFER
+    # =============================================================
+    class ReplayBuffer:
+        def __init__(self, maxlen):
+            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):
+            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
+
+    # =============================================================
+    # PGN
+    # =============================================================
+    def load_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, path, max_games=500):
+        if not os.path.exists(path):
+            return 0, None
+        offset = load_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)
+        log(f"[PGN] loaded chunk={count}, buffer={len(buf):,}")
+        return count, new_offset
+
+    def load_pgn(buf, path, limit=2000):
+        if not os.path.exists(path):
+            log(f"  [PGN] không tìm thấy {path}, bỏ qua.")
+            return
+        count = 0
+        with open(path) as f:
+            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
+        log(f"  [PGN] loaded {count} games, buffer={len(buf):,}")
+
+    # =============================================================
+    # INFERENCE
+    # =============================================================
+    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        = Tensor(states)
+        p_logits, values = model(x)
+        p_logits = p_logits.numpy().astype(np.float64, copy=False)
+        values   = values.numpy()
+        outputs  = []
+        for i, board in enumerate(boards):
+            p = np.nan_to_num(p_logits[i], nan=0.0, posinf=0.0, neginf=0.0, copy=False)
+            mask = np.zeros(ACTION_SIZE, dtype=np.float64)
+            mask[[move_to_idx(m) for m in board.legal_moves]] = 1.0
+            p -= p.max(); np.exp(p, out=p); p *= mask
+            s = p.sum()
+            p = p / s if s > 0 else mask / mask.sum()
+            outputs.append((p.copy(), float(values[i, 0])))
+        return outputs
+
+    # =============================================================
+    # MCTS
+    # =============================================================
+    class MCTSNode:
+        __slots__ = ("board", "prior", "parent", "children", "visits", "total_value")
+        def __init__(self, board, prior=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):
+            return self.total_value / self.visits if self.visits > 0 else 0.0
+        def uct_score(self):
+            pv = self.parent.visits if self.parent else 1
+            return self.q + 1.5 * self.prior * (pv ** 0.5) / (1 + self.visits)
+
+    def terminal_value(board):
+        out = board.outcome()
+        if out is None or out.winner is None: return 0.0
+        return -1.0
+
+    def expand_node(node, policy):
+        for move in node.board.legal_moves:
+            if move not in node.children:
+                cb = node.board.copy(stack=False); cb.push(move)
+                node.children[move] = MCTSNode(cb, prior=float(policy[move_to_idx(move)]), parent=node)
+
+    def backup_path(path, value):
+        for n in reversed(path):
+            n.visits += 1; n.total_value += value; value = -value
+
+    def mcts_search(root, num_sims=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:
+                for idx, (pol, val) in zip(eval_indices, predict_batch(boards_to_eval)):
+                    policies[idx] = pol; values[idx] = val
+
+            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
+
+    # =============================================================
+    # SELF-PLAY
+    # =============================================================
+    def play_one_game(temperature=1.0, game_id=0):
+        board    = chess.Board()
+        history  = []
+        move_num = 0
+        Tensor.training = False
+        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
+        finally:
+            Tensor.training = True
+
+        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 = [(s, p, np.float32(white_z if c == chess.WHITE else -white_z)) for s, p, c in history]
+        return samples, result
+
+    # =============================================================
+    # TRAINING STEP
+    # =============================================================
+    def train_step(buf):
+        if len(buf) < CFG["min_buffer"]:
+            return None
+        loss_policy = loss_value = loss = 0.0
+        skipped = 0
+        for _ in range(CFG["train_steps"]):
+            s, p, z = buf.sample_batch(CFG["batch_size"])
+            S = Tensor(np.ascontiguousarray(s))
+            P = Tensor(np.ascontiguousarray(p))
+            Z = Tensor(np.ascontiguousarray(z))
+            logits, v   = model(S)
+            # clamp logits trước softmax để tránh exp() overflow -> inf
+            logits      = logits.clip(-30, 30)
+            logits      = logits - logits.max(axis=1).reshape(-1, 1)
+            exp_l       = logits.exp()
+            prob        = exp_l / exp_l.sum(axis=1).reshape(-1, 1)
+            loss_policy = -(P * (prob + 1e-8).log()).sum(axis=1).mean()
+            loss_value  = (v - Z).square().mean()
+            loss        = loss_policy + loss_value
+
+            loss_val = float(loss.item())
+            if not np.isfinite(loss_val):
+                # loss NaN/Inf -> bỏ qua step này, KHÔNG cập nhật weights
+                skipped += 1
+                log(f"  [WARN] loss={loss_val} (NaN/Inf) -> skip optimizer.step()")
+                continue
+
+            optimizer.zero_grad(); loss.backward()
+            # gradient clipping: chặn gradient explosion gây NaN lan vào weights
+            for t in get_parameters(model):
+                if t.grad is not None:
+                    t.grad = t.grad.clip(-1.0, 1.0)
+            optimizer.step()
+        return {"loss": float(loss.item()), "loss_policy": float(loss_policy.item()),
+                "loss_value": float(loss_value.item()), "steps": CFG["train_steps"]}
+
+    # =============================================================
+    # MAIN
+    # =============================================================
+    if __name__ == "__main__":
+        step_since_pgn = 0
+        print("=" * 60)
+        print("  Chess AI — AlphaZero-style, tối ưu Colab")
+        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")
+
+        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:
+                import traceback
+                ts = time.strftime("%Y-%m-%d %H:%M:%S")
+                msg = f"[ERROR] self-play: {e}\n{traceback.format_exc()}"
+                log(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:
+                total_steps += 1
+                log(f"  Steps: {total_steps}/{TOTAL_STEPS_TARGET} "
+                    f"loss={info['loss']:.4f} p={info['loss_policy']:.4f} v={info['loss_value']:.4f} "
+                    f"t={elapsed_train:.1f}s")
+
+            loss_str = f"{info['loss']:.4f}" if info else "warmup"
+            log(f"🎮 Game {games_played}: result={result} moves={len(samples)} "
+                f"buf={len(buf):,} loss={loss_str} t={elapsed_game:.1f}s")
+
+            ts   = time.strftime("%Y-%m-%d %H:%M:%S")
+            line = (f"game={games_played:5d} | {result} | moves={len(samples):3d} | buf={len(buf):,} | "
+                    f"steps={total_steps:,} | "
+                    + (f"loss={info['loss']:.4f} (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)
+    return push_to_hf(games_played, total_steps)
+
+
+train()