| |
|
| | PIECE_TO_INDEX = {
|
| | 'wp': 0, 'wN': 1, 'wB': 2, 'wR': 3, 'wQ': 4,
|
| | 'bp': 5, 'bN': 6, 'bB': 7, 'bR': 8, 'bQ': 9
|
| | }
|
| |
|
| |
|
| | NUM_PIECES = 10
|
| | NUM_SQUARES = 64
|
| | NUM_FEATURES = NUM_PIECES * NUM_SQUARES * NUM_SQUARES
|
| |
|
| | PAD_IDX = NUM_FEATURES
|
| | def find_king_squares(board):
|
| | wk = bk = None
|
| | for r in range(8):
|
| | for c in range(8):
|
| | if board[r][c] == "wK":
|
| | wk = r * 8 + c
|
| | elif board[r][c] == "bK":
|
| | bk = r * 8 + c
|
| | return wk, bk
|
| |
|
| | def gs_to_nnue_features(gs):
|
| | board = gs.board
|
| | wk, bk = find_king_squares(board)
|
| |
|
| | features = []
|
| |
|
| | for r in range(8):
|
| | for c in range(8):
|
| | piece = board[r][c]
|
| | if piece == "--" or piece[1] == "K":
|
| | continue
|
| |
|
| | p_idx = PIECE_TO_INDEX[piece]
|
| | sq = r * 8 + c
|
| |
|
| | if piece[0] == 'w':
|
| | king_sq = wk
|
| | else:
|
| | king_sq = bk
|
| |
|
| | if king_sq is None:
|
| | continue
|
| |
|
| |
|
| | idx = (
|
| | p_idx * 64 * 64 +
|
| | king_sq * 64 +
|
| | sq
|
| | )
|
| | features.append(idx)
|
| |
|
| | return features
|
| |
|
| | import torch
|
| |
|
| | class NNUEInfer:
|
| | def __init__(self, model, device="cpu"):
|
| | self.device = device
|
| | self.model = model.to(device)
|
| | self.model.eval()
|
| |
|
| | @torch.no_grad()
|
| | def __call__(self, features, stm):
|
| | """
|
| | features : List[int]
|
| | stm : 1 if white to move, 0 if black
|
| | returns : float score
|
| | """
|
| | if not features:
|
| | features = [PAD_IDX]
|
| |
|
| | feats = torch.tensor(
|
| | features,
|
| | dtype=torch.long,
|
| | device=self.device
|
| | ).unsqueeze(0)
|
| |
|
| | stm = torch.tensor(
|
| | [stm],
|
| | dtype=torch.long,
|
| | device=self.device
|
| | )
|
| |
|
| | return self.model(feats, stm).item()
|
| |
|
