Spaces:
Running
Running
File size: 2,368 Bytes
3f9559b | 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 | #!/usr/bin/env python3
"""Move prediction inference (CNN part) replicating Predict_Human_Move_Train.test_model.
Sanity-checks the 12x8x8 encoding + class order on the start position."""
import sys
import numpy as np
import onnxruntime as ort
import chess
PIECE_ORDER = [chess.PAWN, chess.KNIGHT, chess.BISHOP, chess.ROOK, chess.QUEEN, chess.KING]
SQUARE_MODELS = ["pawn", "knight", "bishop", "rook", "queen", "king"]
def encode(board, row0_rank8=True):
"""12x8x8: channels 0-5 white P,N,B,R,Q,K; 6-11 black. (white-to-move board)"""
enc = np.zeros((12, 8, 8), np.float32)
for sq in chess.SQUARES:
p = board.piece_at(sq)
if p:
c = PIECE_ORDER.index(p.piece_type) + (0 if p.color == chess.WHITE else 6)
r = 7 - (sq // 8) if row0_rank8 else (sq // 8)
enc[c, r, sq % 8] = 1.0
return enc[None]
def softmax(x):
e = np.exp(x - x.max()); return e / e.sum()
def cnn_moves(board, piece_sess, square_sesss, row0_rank8=True):
enc = encode(board, row0_rank8)
pieces = softmax(piece_sess.run(None, {"input": enc})[0].flatten())
legal = list(board.legal_moves)
move_prob = {}
for i, ptype in enumerate(PIECE_ORDER):
squares = softmax(square_sesss[i].run(None, {"input": enc})[0].flatten())
squares = (squares + pieces[i]) / 2
for from_sq in board.pieces(ptype, chess.WHITE):
fs = chess.square_name(from_sq)
for j in range(64):
try:
mv = chess.Move.from_uci(fs + chess.square_name(j))
if mv in legal:
move_prob[mv.uci()] = squares[j]
except Exception:
pass
return [m for m, _ in sorted(move_prob.items(), key=lambda x: x[1], reverse=True)]
def main():
piece_sess = ort.InferenceSession("models_onnx/piece.int8.onnx", providers=["CPUExecutionProvider"])
square_sesss = [ort.InferenceSession(f"models_onnx/square_{s}.int8.onnx",
providers=["CPUExecutionProvider"]) for s in SQUARE_MODELS]
board = chess.Board() # start position (white to move)
for orient in (True, False):
moves = cnn_moves(board, piece_sess, square_sesss, row0_rank8=orient)
print(f"row0_rank8={orient}: top CNN moves -> {moves[:8]}")
if __name__ == "__main__":
main()
|