visualize nnue input weights

nnue_visualizer.py

from typing import Dict
import json, sys
from PIL import Image

def load_state(path) -> Dict[str, list]:
    with open(path) as f:
        return json.load(f)

state = load_state(sys.argv[1])

ft_w = state["weights"]

max_v = 255
min_v = -255

il = 768
hl = 128

col = 8
row = 16

w = col * (8 * 6 + 1) + 1
h = row * (8 * 2 + 1) + 1

image = Image.new("RGB", (w, h), (255, 255, 255))

def feature(color, piece, square):
    return 384 * color + 64 * piece + square

max_hl = 0

for hl_idx, weights in enumerate(ft_w):
    base_x = int(hl_idx % col) * (8 * 6 + 1) + 1
    base_y = int(hl_idx / col) * (8 * 2 + 1) + 1

    for color in range(2):
        for piece in range(6):
            for x in range(8):
                for y in range(8):
                    sq = ((7 - y) << 3) | x

                    w = weights[feature(color, piece, sq)]

                    px = abs(w)
                    px = [(px, 0, 0), (0, px, 0)][w >= 0]

                    image.putpixel((base_x + piece * 8 + x, base_y + color * 8 + y), px)

image.save("out.png")