Compatibility with resnet architecture (#8)

* old resnet archi

* new onnx cannonisation

* increased compatibility

* make datasets

* fixed tests

* sample exploration

Author: Xmaster6y

.vscode/launch.json

@@ -19,20 +19,37 @@
             "justMyCode": false
         },
         {
-            "name": "Script CRP Concepts",
+            "name": "Script CRP concepts",
             "type": "debugpy",
             "request": "launch",
             "module": "scripts.find_concepts",
             "console": "integratedTerminal",
             "justMyCode": false
         },
         {
-            "name": "Script CRP Clusters",
+            "name": "Script CRP clusters",
             "type": "debugpy",
             "request": "launch",
             "module": "scripts.cluster_latent_relevances",
             "console": "integratedTerminal",
             "justMyCode": false
+        },
+        {
+            "name": "Script make datasets",
+            "type": "debugpy",
+            "request": "launch",
+            "module": "scripts.make_datasets",
+            "console": "integratedTerminal",
+            "justMyCode": false
+        }
+        ,
+        {
+            "name": "Script sample exploration",
+            "type": "debugpy",
+            "request": "launch",
+            "module": "scripts.sample_exploration",
+            "console": "integratedTerminal",
+            "justMyCode": false
         }
     ]
 }

demo/policy_interface.py

@@ -67,7 +67,7 @@ def compute_policy(
             gr.Warning("Invalid action sequence.")
             return (None, None, "", None)
     wrapper = utils.get_wrapper_from_state(model_name)
-    output = wrapper.predict(board)
+    (output,) = wrapper.predict(board)
     current_raw_policy = output["policy"][0]
     policy = torch.softmax(output["policy"][0], dim=-1)
 

demo/statistics_interface.py

@@ -6,16 +6,15 @@
 
 from demo import utils, visualisation
 from lczerolens import GameDataset
-from lczerolens.xai import HasThreatConcept, UniqueConceptDataset
+from lczerolens.xai import ConceptDataset, HasThreatConcept
 
 current_policy_statistics = None
 current_lrp_statistics = None
 current_probing_statistics = None
 dataset = GameDataset("assets/test_stockfish_10.jsonl")
 check_concept = HasThreatConcept("K", relative=True)
-unique_check_dataset = UniqueConceptDataset.from_game_dataset(
-    dataset, check_concept
-)
+unique_check_dataset = ConceptDataset.from_game_dataset(dataset)
+unique_check_dataset.concept = check_concept
 
 
 def list_models():
@@ -47,7 +46,7 @@ def compute_policy_statistics(
         )
         return None
     wrapper, lens = utils.get_wrapper_lens_from_state(model_name, "policy")
-    current_policy_statistics = lens.compute_statistics(dataset, wrapper, 10)
+    current_policy_statistics = lens.analyse_dataset(dataset, wrapper, 10)
     return make_policy_plot()
 
 

demo/utils.py

@@ -7,7 +7,7 @@
 import subprocess
 
 from demo import constants, state
-from lczerolens import AutoLens, ModelWrapper
+from lczerolens import Lens, ModelWrapper
 from lczerolens.utils import lczero as lczero_utils
 
 
@@ -122,7 +122,7 @@ def get_wrapper_lens_from_state(
     if lens_name in state.lenses[lens_type]:
         lens = state.lenses[lens_type][lens_name]
     else:
-        lens = AutoLens.from_type(lens_type, **kwargs)
+        lens = Lens.from_name(lens_type, **kwargs)
         if not lens.is_compatible(wrapper):
             raise ValueError(
                 f"Lens of type {lens_type} not compatible with model."

demo/visualisation.py

@@ -102,7 +102,7 @@ def render_architecture(model, name: str = "model", directory: str = ""):
 def render_policy_distribution(
     policy,
     legal_moves,
-    n_bins=10,
+    n_bins=20,
 ):
     """
     Render the policy distribution histogram.
@@ -112,11 +112,12 @@ def render_policy_distribution(
     ).bool()
     fig = plt.figure(figsize=(6, 6))
     ax = plt.gca()
-    _, bins, _ = ax.hist(
+    _, bins = np.histogram(policy, bins=n_bins)
+    ax.hist(
         policy[~legal_mask],
-        bins=n_bins,
-        density=True,
+        bins=bins,
         alpha=0.5,
+        density=True,
         label="Illegal moves",
     )
     ax.hist(

scripts/cluster_latent_relevances.py

@@ -14,7 +14,6 @@
 import torch
 from crp.attribution import CondAttribution
 from crp.concepts import ChannelConcept
-from crp.helper import get_layer_names
 from pylatex import Document
 from pylatex.package import Package
 from safetensors import safe_open
@@ -23,88 +22,104 @@
 from sklearn.manifold import TSNE
 from tqdm import tqdm
 
-from lczerolens import GameDataset, move_utils
-from lczerolens.adapt import PolicyFlow
-from lczerolens.xai import LrpLens, UniqueConceptDataset
-from lczerolens.xai.concepts import BestLegalMoveConcept
+from lczerolens import move_utils
+from lczerolens.game import PolicyFlow
+from lczerolens.xai import ConceptDataset, LrpLens
 from scripts.create_figure import add_plot, create_heatmap_string
 
 #######################################
 # HYPERPARAMETERS
 #######################################
 n_clusters = 10
-layer_index = -1
 batch_size = 500
-save_files = False
+save_files = True
 conv_sum_dims = (2, 3)
-model_name = "tinygyal-8.onnx"
-dataset_name = "test_stockfish_10.jsonl"
+model_name = "64x6-2018_0627_1913_08_161.onnx"
+dataset_name = "TCEC_game_collection_random_boards_bestlegal_knight.jsonl"
 only_config_rel = True
+best_legal = True
+run_name = (
+    f"bestres_tcec_bestlegal_knight_{'expbest' if best_legal else 'full'}"
+)
 #######################################
 
 
-class MaxLogitFlow(PolicyFlow):
-    def forward(self, x):
-        policy = super().forward(x)
-        return policy.max(dim=1, keepdim=True).values
+def legal_init_rel(board_list, board_tensor):
+    legal_move_mask = torch.zeros((len(board_list), 1858))
+    for idx, board in enumerate(board_list):
+        legal_moves = [
+            move_utils.encode_move(move, (board.turn, not board.turn))
+            for move in board.legal_moves
+        ]
+        legal_move_mask[idx, legal_moves] = 1
+    return legal_move_mask * board_tensor
 
 
-model = MaxLogitFlow.from_path(f"./assets/{model_name}")
-dataset = GameDataset(f"./assets/{dataset_name}")
-concept = BestLegalMoveConcept(model)
-unique_dataset = UniqueConceptDataset.from_game_dataset(dataset, concept)
-print(f"[INFO] Board dataset len: {len(unique_dataset)}")
+model = PolicyFlow.from_path(f"./assets/{model_name}")
+concept_dataset = ConceptDataset(f"./assets/{dataset_name}")
+print(f"[INFO] Board dataset len: {len(concept_dataset)}")
 
 composite = LrpLens.make_default_composite()
-attribution = CondAttribution(model)
 cc = ChannelConcept()
-
-layer_names = get_layer_names(model, [torch.nn.ReLU])
-layer_names = [
-    layer_name for layer_name in layer_names if "block" in layer_name
-]
+layer_names = [f"model.block{b}/conv2/relu" for b in [0, 3, 5]]
 print(layer_names)
 
 dataloader = torch.utils.data.DataLoader(
-    unique_dataset,
+    concept_dataset,
     batch_size=batch_size,
     shuffle=False,
-    collate_fn=UniqueConceptDataset.collate_fn_tensor,
+    collate_fn=ConceptDataset.collate_fn_tensor,
 )
 
 if save_files:
     print("############ Collecting Relevances")
     all_relevances = {}
     for batch in tqdm(dataloader):
-        _, board_tensor, _ = batch
+        _, board_tensor, labels = batch
+        label_tensor = torch.tensor(labels)
+
+        def init_rel_fn(board_tensor):
+            rel = torch.zeros_like(board_tensor)
+            rel[:, label_tensor] = board_tensor[:, label_tensor]
+            return rel
+
         board_tensor.requires_grad = True
-        attr = attribution(
-            board_tensor, [{"y": 0}], composite, record_layer=layer_names
-        )
-
-        for layer_name in layer_names:
-            latent_rel = attr.relevances[layer_name]
-            latent_rel = cc.attribute(latent_rel, abs_norm=True)
-            if len(latent_rel.shape) == 4:
-                latent_rel = latent_rel.sum(conv_sum_dims)
-            if layer_name not in all_relevances:
-                all_relevances[layer_name] = latent_rel.detach().cpu()
-            else:
-                all_relevances[layer_name] = torch.cat(
-                    [all_relevances[layer_name], latent_rel.detach().cpu()],
-                    dim=0,
-                )
+        with LrpLens.context(model) as modifed_model:
+            attribution = CondAttribution(modifed_model)
+            attr = attribution(
+                board_tensor,
+                [{"y": None}],
+                composite,
+                record_layer=layer_names,
+                init_rel=init_rel_fn if best_legal else None,
+            )
 
-    os.makedirs(f"scripts/clusters/{model_name}-{dataset_name}", exist_ok=True)
+            for layer_name in layer_names:
+                latent_rel = attr.relevances[layer_name]
+                latent_rel = cc.attribute(latent_rel, abs_norm=True)
+                if len(latent_rel.shape) == 4:
+                    latent_rel = latent_rel.sum(conv_sum_dims)
+                if layer_name not in all_relevances:
+                    all_relevances[layer_name] = latent_rel.detach().cpu()
+                else:
+                    all_relevances[layer_name] = torch.cat(
+                        [
+                            all_relevances[layer_name],
+                            latent_rel.detach().cpu(),
+                        ],
+                        dim=0,
+                    )
+
+    os.makedirs(f"scripts/clusters/{run_name}", exist_ok=True)
     save_file(
         all_relevances,
-        f"scripts/clusters/{model_name}-{dataset_name}/relevances.safetensors",
+        f"scripts/clusters/{run_name}/relevances.safetensors",
     )
 
 else:
     all_relevances = {}
     with safe_open(
-        f"scripts/clusters/{model_name}-{dataset_name}/relevances.safetensors",
+        f"scripts/clusters/{run_name}/relevances.safetensors",
         framework="pt",
         device="cpu",
     ) as f:
@@ -116,7 +131,7 @@ def forward(self, x):
 #######################################
 
 print("############ Clustering ...")
-os.makedirs(f"scripts/results/{model_name}-{dataset_name}", exist_ok=True)
+os.makedirs(f"scripts/results/{run_name}", exist_ok=True)
 
 for layer_name, relevances in all_relevances.items():
     kmeans = KMeans(n_clusters=n_clusters, init="k-means++")
@@ -132,7 +147,7 @@ def forward(self, x):
     plt.xlabel("Dimension 1")
     plt.ylabel("Dimension 2")
     plt.savefig(
-        f"scripts/results/{model_name}-{dataset_name}/{layer_name}_t-sne.png"
+        f"scripts/results/{run_name}/{layer_name.replace('/','.')}_t-sne.png"
     )
     plt.close()
 
@@ -141,47 +156,60 @@ def forward(self, x):
     #######################################
 
     print("############ Plotting chessboards for each cluster")
+    with LrpLens.context(model) as modifed_model:
+        attribution = CondAttribution(modifed_model)
+        for idx_cluster in tqdm(range(n_clusters)):
+            cluster_center = kmeans.cluster_centers_[idx_cluster]
+            distances = np.linalg.norm(relevances - cluster_center, axis=1)
+            nearest_neighbors = np.argsort(distances)[:8]
+
+            doc = Document()  # create a new document
+            doc.packages.append(Package("xskak"))
+
+            # compute heatmap for each nearest neighbor
+            for idx_sample in nearest_neighbors:
+                _, board, label = concept_dataset[idx_sample]
+                _, board_tensor, _ = ConceptDataset.collate_fn_tensor(
+                    [concept_dataset[idx_sample]]
+                )
+                label_tensor = torch.tensor([label])
+
+                def init_rel_fn(board_tensor):
+                    rel = torch.zeros_like(board_tensor)
+                    rel[:, label_tensor] = board_tensor[:, label_tensor]
+                    return rel
+
+                board_tensor.requires_grad = True
+                attr = attribution(
+                    board_tensor,
+                    [{"y": None}],
+                    composite,
+                    init_rel=init_rel_fn if best_legal else None,
+                )
+                if only_config_rel:
+                    heatmap = board_tensor.grad[0, :12].sum(dim=0).view(64)
+                else:
+                    heatmap = board_tensor.grad[0].sum(dim=0).view(64)
+                if board.turn == chess.BLACK:
+                    heatmap = heatmap.view(8, 8).flip(0).view(64)
+                move = move_utils.decode_move(
+                    label, (board.turn, not board.turn), board
+                )
+                uci_move = move.uci()
+                heatmap = heatmap / heatmap.abs().max()
+                heatmap_str = create_heatmap_string(heatmap)
+
+                doc = add_plot(
+                    doc,
+                    board.fen(),
+                    heatmap_str,
+                    current_piece_pos=uci_move[:2],
+                    next_move=uci_move[2:4],
+                )
 
-    for idx_cluster in tqdm(range(n_clusters)):
-        cluster_center = kmeans.cluster_centers_[idx_cluster]
-        distances = np.linalg.norm(relevances - cluster_center, axis=1)
-        nearest_neighbors = np.argsort(distances)[:10]
-
-        doc = Document()  # create a new document
-        doc.packages.append(Package("xskak"))
-
-        # compute heatmap for each nearest neighbor
-        for idx_sample in nearest_neighbors:
-            _, board, label = unique_dataset[idx_sample]
-            _, board_tensor, _ = UniqueConceptDataset.collate_fn_tensor(
-                [unique_dataset[idx_sample]]
+            # Generate pdf
+            doc.generate_pdf(
+                f"scripts/results/{run_name}"
+                f"/{layer_name.replace('/','.')}_cluster_{idx_cluster}",
+                clean_tex=True,
             )
-            board_tensor.requires_grad = True
-            attr = attribution(board_tensor, [{"y": 0}], composite)
-            if only_config_rel:
-                heatmap = board_tensor.grad[0, :12].sum(dim=0).view(64)
-            else:
-                heatmap = board_tensor.grad[0].sum(dim=0).view(64)
-            if board.turn == chess.BLACK:
-                heatmap = heatmap.view(8, 8).flip(0).view(64)
-            move = move_utils.decode_move(
-                label, (board.turn, not board.turn), board
-            )
-            uci_move = move.uci()
-            heatmap = heatmap / heatmap.abs().max()
-            heatmap_str = create_heatmap_string(heatmap)
-
-            doc = add_plot(
-                doc,
-                board.fen(),
-                heatmap_str,
-                current_piece_pos=uci_move[:2],
-                next_move=uci_move[2:4],
-            )
-
-        # Generate pdf
-        doc.generate_pdf(
-            f"scripts/results/{model_name}-{dataset_name}"
-            f"/{layer_name}_cluster_{idx_cluster}",
-            clean_tex=True,
-        )