input_encoding

demo/interfaces/activations.py

@@ -9,7 +9,7 @@ import gradio as gr
 import os
 import torch
 
-from lczerolens import LczeroBoard, LczeroModel, Lens
+from lczerolens import LczeroBoard, LczeroModel, Lens, InputEncoding
 
 from demo import constants
 from demo.utils import get_info
@@ -17,10 +17,10 @@ from demo.utils import get_info
 def get_model(model_name: str):
     return LczeroModel.from_onnx_path(os.path.join(constants.ONNX_MODEL_DIRECTORY, model_name))
 
-def get_activations(model: LczeroModel, board: LczeroBoard):
+def get_activations(model: LczeroModel, board: LczeroBoard, input_encoding: InputEncoding):
     lens = Lens.from_name("activation", "block\d/conv2/relu")
     with torch.no_grad():
-        results = lens.analyse(model, board)
+        results = lens.analyse(model, board, model_kwargs={"input_encoding": input_encoding})
     return [results[f"block{i}/conv2/relu_output"][0] for i in range(len(results))]
 
 def get_board(game_pgn:str, board_fen:str):
@@ -62,28 +62,34 @@ def render_activations(board: LczeroBoard, activations, layer_index:int, channel
     )
     return  f"{constants.FIGURE_DIRECTORY}/activations_board.svg", f"{constants.FIGURE_DIRECTORY}/activations_colorbar.svg"
 
-def initial_load(model_name: str, board_fen: str, game_pgn: str, layer_index: int, channel_index: int):
+def initial_load(model_name: str, board_fen: str, game_pgn: str, input_encoding: InputEncoding, layer_index: int, channel_index: int):
     model = get_model(model_name)
     board = get_board(game_pgn, board_fen)
-    activations = get_activations(model, board)
+    activations = get_activations(model, board, input_encoding)
     info = get_info(model, board)
     plots = render_activations(board, activations, layer_index, channel_index)
     return model, board, activations, info, *plots
 
-def on_board_change(model: LczeroModel, game_pgn: str, board_fen: str, layer_index: int, channel_index: int):
+def on_board_change(model: LczeroModel, game_pgn: str, board_fen: str, input_encoding: InputEncoding, layer_index: int, channel_index: int):
     board = get_board(game_pgn, board_fen)
-    activations = get_activations(model, board)
+    activations = get_activations(model, board, input_encoding)
     info = get_info(model, board)
     plots = render_activations(board, activations, layer_index, channel_index)
     return board, activations, info, *plots
 
-def on_model_change(model_name: str, board: LczeroBoard, layer_index: int, channel_index: int):
+def on_model_change(model_name: str, board: LczeroBoard, input_encoding: InputEncoding, layer_index: int, channel_index: int):
     model = get_model(model_name)
-    activations = get_activations(model, board)
+    activations = get_activations(model, board, input_encoding)
     info = get_info(model, board)
     plots = render_activations(board, activations, layer_index, channel_index)
     return model, activations, info, *plots
 
+def on_input_encoding_change(model: LczeroModel, board: LczeroBoard, input_encoding: InputEncoding, layer_index: int, channel_index: int):
+    activations = get_activations(model, board, input_encoding)
+    info = get_info(model, board)
+    plots = render_activations(board, activations, layer_index, channel_index)
+    return activations, info, *plots
+
 with gr.Blocks() as interface:
     with gr.Row():
         with gr.Column():
@@ -102,6 +108,15 @@ with gr.Blocks() as interface:
                     max_lines=1,
                     value=chess.STARTING_FEN,
                 )
+                input_encoding = gr.Radio(
+                    label="Input encoding",
+                    choices=[
+                        ("classical", InputEncoding.INPUT_CLASSICAL_112_PLANE), 
+                        ("repeated", InputEncoding.INPUT_CLASSICAL_112_PLANE_REPEATED), 
+                        ("no history", InputEncoding.INPUT_CLASSICAL_112_PLANE_NO_HISTORY)
+                    ],
+                    value=InputEncoding.INPUT_CLASSICAL_112_PLANE,
+                )
                 model_name = gr.Dropdown(
                     label="Model",
                     choices=constants.ONNX_MODEL_NAMES,
@@ -133,36 +148,42 @@ with gr.Blocks() as interface:
 
     interface.load(
         initial_load,
-        inputs=[model_name, game_pgn, board_fen, layer_index, channel_index],
+        inputs=[model_name, game_pgn, board_fen, input_encoding, layer_index, channel_index],
         outputs=[model, board, activations, info, image_board, colorbar],
         concurrency_limit=1, 
         concurrency_id="trace_queue"
     )
     game_pgn.submit(
         on_board_change,
-        inputs=[model, game_pgn, board_fen, layer_index, channel_index],
+        inputs=[model, game_pgn, board_fen, input_encoding, layer_index, channel_index],
         outputs=[board, activations, info, image_board, colorbar],
         concurrency_id="trace_queue"
     )
     board_fen.submit(
         on_board_change,
-        inputs=[model, game_pgn, board_fen, layer_index, channel_index],
+        inputs=[model, game_pgn, board_fen, input_encoding, layer_index, channel_index],
         outputs=[board, activations, info, image_board, colorbar],
         concurrency_id="trace_queue"
     )
     model_name.change(
         on_model_change,
-        inputs=[model_name, board, layer_index, channel_index],
+        inputs=[model_name, board, input_encoding, layer_index, channel_index],
         outputs=[model, activations, info, image_board, colorbar],
         concurrency_id="trace_queue"
     )
+    input_encoding.change(
+        on_input_encoding_change,
+        inputs=[model, board, input_encoding, layer_index, channel_index],
+        outputs=[activations, info, image_board, colorbar],
+        concurrency_id="trace_queue"
+    )
     layer_index.change(
         render_activations,
-        inputs=[board, activations, layer_index, channel_index],
+        inputs=[board, activations, input_encoding, layer_index, channel_index],
         outputs=[image_board, colorbar],
     )
     channel_index.change(
         render_activations,
-        inputs=[board, activations, layer_index, channel_index],
+        inputs=[board, activations, input_encoding, layer_index, channel_index],
         outputs=[image_board, colorbar],
     )

demo/interfaces/encodings.py

@@ -7,11 +7,11 @@ import chess.pgn
 import io
 import gradio as gr
 
-from lczerolens.board import LczeroBoard
+from lczerolens.board import LczeroBoard, InputEncoding
 
 from ..constants import FIGURE_DIRECTORY
 
-def make_render(game_pgn:str, board_fen:str, plane_index:int):
+def make_render(game_pgn:str, board_fen:str, input_encoding:InputEncoding, plane_index:int):
     if game_pgn:
         try:
             board = LczeroBoard()
@@ -30,10 +30,10 @@ def make_render(game_pgn:str, board_fen:str, plane_index:int):
             print(e)
             gr.Warning("Invalid FEN, using starting position.")
             board = LczeroBoard()
-    return board, *make_board_plot(board, plane_index)
+    return board, *make_board_plot(board, input_encoding, plane_index)
 
-def make_board_plot(board:LczeroBoard, plane_index:int):
-    input_tensor = board.to_input_tensor()
+def make_board_plot(board:LczeroBoard, input_encoding:InputEncoding, plane_index:int):
+    input_tensor = board.to_input_tensor(input_encoding)
     board.render_heatmap(
         input_tensor[plane_index].view(64),
         save_to=f"{FIGURE_DIRECTORY}/encodings.svg",
@@ -60,6 +60,15 @@ with gr.Blocks() as interface:
                     max_lines=1,
                     value=chess.STARTING_FEN,
                 )
+                input_encoding = gr.Radio(
+                    label="Input encoding",
+                    choices=[
+                        ("classical", InputEncoding.INPUT_CLASSICAL_112_PLANE), 
+                        ("repeated", InputEncoding.INPUT_CLASSICAL_112_PLANE_REPEATED), 
+                        ("no history", InputEncoding.INPUT_CLASSICAL_112_PLANE_NO_HISTORY)
+                    ],
+                    value=InputEncoding.INPUT_CLASSICAL_112_PLANE,
+                )
             with gr.Group():
                 with gr.Row():
                     plane_index = gr.Slider(
@@ -75,7 +84,7 @@ with gr.Blocks() as interface:
 
     state_board = gr.State(value=LczeroBoard())
 
-    render_inputs = [game_pgn, board_fen, plane_index]
+    render_inputs = [game_pgn, board_fen, input_encoding, plane_index]
     render_outputs = [state_board, image_board, colorbar]
     interface.load(
         make_render,
@@ -92,8 +101,13 @@ with gr.Blocks() as interface:
         inputs=render_inputs,
         outputs=render_outputs,
     )
+    input_encoding.change(
+        make_board_plot,
+        inputs=[state_board, input_encoding, plane_index],
+        outputs=[image_board, colorbar],
+    )
     plane_index.change(
         make_board_plot,
-        inputs=[state_board, plane_index],
+        inputs=[state_board, input_encoding, plane_index],
         outputs=[image_board, colorbar],
     )

demo/interfaces/gradients.py

@@ -8,7 +8,7 @@ import io
 import gradio as gr
 import os
 
-from lczerolens import LczeroBoard, LczeroModel, Lens
+from lczerolens import LczeroBoard, LczeroModel, Lens, InputEncoding
 
 from demo import constants
 from demo.utils import get_info
@@ -16,7 +16,7 @@ from demo.utils import get_info
 def get_model(model_name: str):
     return LczeroModel.from_onnx_path(os.path.join(constants.ONNX_MODEL_DIRECTORY, model_name))
 
-def get_gradients(model: LczeroModel, board: LczeroBoard, target: str):
+def get_gradients(model: LczeroModel, board: LczeroBoard, input_encoding: InputEncoding, target: str):
     lens = Lens.from_name("gradient")
 
     def init_target(model):
@@ -25,7 +25,7 @@ def get_gradients(model: LczeroModel, board: LczeroBoard, target: str):
         else:
             wdl_index = {"win": 0, "draw": 1, "loss": 2}[target]
             return getattr(model, "output/wdl").output[:, wdl_index]
-    results = lens.analyse(model, board, init_target=init_target)
+    results = lens.analyse(model, board, init_target=init_target, model_kwargs={"input_encoding": input_encoding})
 
     return results["input_grad"]
 
@@ -61,30 +61,35 @@ def render_gradients(board: LczeroBoard, gradients, average_over_planes:bool, be
     )
     return  f"{constants.FIGURE_DIRECTORY}/gradients_board.svg", f"{constants.FIGURE_DIRECTORY}/gradients_colorbar.svg"
 
-def initial_load(model_name: str, board_fen: str, game_pgn: str, target: str, average_over_planes:bool, begin_average_index:int, end_average_index:int, plane_index: int):
+def initial_load(model_name: str, board_fen: str, game_pgn: str, input_encoding: InputEncoding, target: str, average_over_planes:bool, begin_average_index:int, end_average_index:int, plane_index: int):
     model = get_model(model_name)
     board = get_board(game_pgn, board_fen)
-    gradients = get_gradients(model, board, target)
+    gradients = get_gradients(model, board, input_encoding, target)
     info = get_info(model, board)
     plots = render_gradients(board, gradients, average_over_planes, begin_average_index, end_average_index, plane_index)
     return model, board, gradients, info, *plots
 
-def on_board_change(model: LczeroModel, game_pgn: str, board_fen: str, target: str, average_over_planes:bool, begin_average_index:int, end_average_index:int, plane_index: int):
+def on_board_change(model: LczeroModel, game_pgn: str, board_fen: str, input_encoding: InputEncoding, target: str, average_over_planes:bool, begin_average_index:int, end_average_index:int, plane_index: int):
     board = get_board(game_pgn, board_fen)
-    gradients = get_gradients(model, board, target)
+    gradients = get_gradients(model, board, input_encoding, target)
     info = get_info(model, board)
     plots = render_gradients(board, gradients, average_over_planes, begin_average_index, end_average_index, plane_index)
     return board, gradients, info, *plots
 
-def on_model_change(model_name: str, board: LczeroBoard, target: str, average_over_planes:bool, begin_average_index:int, end_average_index:int, plane_index: int):
+def on_model_change(model_name: str, board: LczeroBoard, input_encoding: InputEncoding, target: str, average_over_planes:bool, begin_average_index:int, end_average_index:int, plane_index: int):
     model = get_model(model_name)
-    gradients = get_gradients(model, board, target)
+    gradients = get_gradients(model, board, input_encoding, target)
     info = get_info(model, board)
     plots = render_gradients(board, gradients, average_over_planes, begin_average_index, end_average_index, plane_index)
     return model, gradients, info, *plots
 
-def on_target_change(model: LczeroModel, board: LczeroBoard, target: str, average_over_planes:bool, begin_average_index:int, end_average_index:int, plane_index: int):
-    gradients = get_gradients(model, board, target)
+def on_input_encoding_change(model: LczeroModel, board: LczeroBoard, input_encoding: InputEncoding, target: str, average_over_planes:bool, begin_average_index:int, end_average_index:int, plane_index: int):
+    gradients = get_gradients(model, board, input_encoding, target)
+    plots = render_gradients(board, gradients, average_over_planes, begin_average_index, end_average_index, plane_index)
+    return gradients, *plots
+
+def on_target_change(model: LczeroModel, board: LczeroBoard, input_encoding: InputEncoding, target: str, average_over_planes:bool, begin_average_index:int, end_average_index:int, plane_index: int):
+    gradients = get_gradients(model, board, input_encoding, target)
     plots = render_gradients(board, gradients, average_over_planes, begin_average_index, end_average_index, plane_index)
     return gradients, *plots
 
@@ -106,6 +111,15 @@ with gr.Blocks() as interface:
                     max_lines=1,
                     value=chess.STARTING_FEN,
                 )
+                input_encoding = gr.Radio(
+                    label="Input encoding",
+                    choices=[
+                        ("classical", InputEncoding.INPUT_CLASSICAL_112_PLANE), 
+                        ("repeated", InputEncoding.INPUT_CLASSICAL_112_PLANE_REPEATED), 
+                        ("no history", InputEncoding.INPUT_CLASSICAL_112_PLANE_NO_HISTORY)
+                    ],
+                    value=InputEncoding.INPUT_CLASSICAL_112_PLANE,
+                )
                 model_name = gr.Dropdown(
                     label="Model",
                     choices=constants.ONNX_MODEL_NAMES,
@@ -150,37 +164,43 @@ with gr.Blocks() as interface:
 
     interface.load(
         initial_load,
-        inputs=[model_name, game_pgn, board_fen, target, average_over_planes, begin_average_index, end_average_index, plane_index],
+        inputs=[model_name, game_pgn, board_fen, input_encoding, target, average_over_planes, begin_average_index, end_average_index, plane_index],
         outputs=[model, board, gradients, info, image_board, colorbar],
         concurrency_id="trace_queue"
     )
     game_pgn.submit(
         on_board_change,
-        inputs=[model, game_pgn, board_fen, target, average_over_planes, begin_average_index, end_average_index, plane_index],
+        inputs=[model, game_pgn, board_fen, input_encoding, target, average_over_planes, begin_average_index, end_average_index, plane_index],
         outputs=[board, gradients, info, image_board, colorbar],
         concurrency_id="trace_queue"
     )
     board_fen.submit(
         on_board_change,
-        inputs=[model, game_pgn, board_fen, target, average_over_planes, begin_average_index, end_average_index, plane_index],
+        inputs=[model, game_pgn, board_fen, input_encoding, target, average_over_planes, begin_average_index, end_average_index, plane_index],
         outputs=[board, gradients, info, image_board, colorbar],
         concurrency_id="trace_queue"
     )
     model_name.change(
         on_model_change,
-        inputs=[model_name, board, target, average_over_planes, begin_average_index, end_average_index, plane_index],
+        inputs=[model_name, board, input_encoding, target, average_over_planes, begin_average_index, end_average_index, plane_index],
         outputs=[model, gradients, info, image_board, colorbar],
         concurrency_id="trace_queue"
     )
+    input_encoding.change(
+        on_input_encoding_change,
+        inputs=[model, board, input_encoding, target, average_over_planes, begin_average_index, end_average_index, plane_index],
+        outputs=[gradients, image_board, colorbar],
+        concurrency_id="trace_queue"
+    )
     target.change(
         on_target_change,
-        inputs=[model, board, target, average_over_planes, begin_average_index, end_average_index, plane_index],
+        inputs=[model, board, input_encoding, target, average_over_planes, begin_average_index, end_average_index, plane_index],
         outputs=[gradients, image_board, colorbar],
         concurrency_id="trace_queue"
     )
     for render_arg in [average_over_planes, begin_average_index, end_average_index, plane_index]:
         render_arg.change(
             render_gradients,
-            inputs=[board, gradients, average_over_planes, begin_average_index, end_average_index, plane_index],
+            inputs=[board, gradients, input_encoding, average_over_planes, begin_average_index, end_average_index, plane_index],
             outputs=[image_board, colorbar],
         )