Extended clustering script and apptainer (#11)

* expanded clustering

* running scripts with apptainer

* gpu label

Author: Xmaster6y

apptainer/.gitignore

@@ -0,0 +1,5 @@
+*
+!.gitignore
+!base.def
+!script.def
+!make-datasets.sh

apptainer/base.def

@@ -0,0 +1,19 @@
+Bootstrap: docker
+From: python:3.9.18
+
+%files
+  ./assets /opt/assets
+  ./src /opt/src
+  ./pyproject.toml /opt/pyproject.toml
+  ./poetry.lock /opt/poetry.lock
+  ./README.md /opt/README.md
+
+%environment
+  export "PATH=/opt/.venv/bin:$PATH"
+
+%post
+  python -m pip install poetry
+
+  cd /opt
+  python -m poetry config virtualenvs.in-project true
+  python -m poetry install

apptainer/make-datasets.sh

@@ -0,0 +1,29 @@
+#!/bin/bash
+
+#SBATCH --mail-type=ALL
+#SBATCH --mail-user=<[email protected]>
+#SBATCH --job-name=apptainer
+#SBATCH --output=%j_%x.out
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=4
+#SBATCH --gpus=1
+#SBATCH --mem=32G
+#SBATCH --time=1:00:00
+
+#####################################################################################
+
+# This included file contains the definition for $LOCAL_JOB_DIR to be used locally on the node.
+source "/etc/slurm/local_job_dir.sh"
+
+# Launch the apptainer image with --nv for nvidia support. Two bind mounts are used:
+# - One for the ImageNet dataset and
+# - One for the results (e.g. checkpoint data that you may store in $LOCAL_JOB_DIR on the node
+timeout 24h apptainer exec --nv --bind ${LOCAL_JOB_DIR}:/opt/output \
+    ./apptainer/script.sif python -m scripts.make_datasets \
+    --output-root /opt/output
+
+# This command copies all results generated in $LOCAL_JOB_DIR back to the submit folder regarding the job id.
+cp -r ${LOCAL_JOB_DIR} ${SLURM_SUBMIT_DIR}/${SLURM_JOB_ID}
+
+echo "$PWD/${SLURM_JOB_ID}_stats.out" > $LOCAL_JOB_DIR/stats_file_loc_cfg

apptainer/script.def

@@ -0,0 +1,9 @@
+Bootstrap: localimage
+From: ./apptainer/base.sif
+
+%files
+  ./scripts/*.py /opt/scripts/
+
+%runscript
+  cd /opt/
+  echo "Running script"

scripts/cluster_latent_relevances.py

@@ -8,12 +8,10 @@
 
 import os
 
-import chess
 import matplotlib.pyplot as plt
 import numpy as np
 import torch
 from crp.attribution import CondAttribution
-from crp.concepts import ChannelConcept
 from pylatex import Document
 from pylatex.package import Package
 from safetensors import safe_open
@@ -30,17 +28,26 @@
 #######################################
 # HYPERPARAMETERS
 #######################################
-n_clusters = 10
+n_clusters = 15
 batch_size = 500
-save_files = True
-conv_sum_dims = (2, 3)
+save_files = False
 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'}"
 )
+n_samples = 1000
+conv_sum_dims = ()
+cosine_sim = False
+kmeans_on_tsne = True
+viz_latent = True
+viz_name = (
+    f"{'latent' if viz_latent else 'input'}"
+    f"_nosum_{'cosine' if cosine_sim else 'norm'}"
+    f"_{'after' if kmeans_on_tsne else 'before'}-tsne"
+)
 #######################################
 
 
@@ -60,7 +67,6 @@ def legal_init_rel(board_list, board_tensor):
 print(f"[INFO] Board dataset len: {len(concept_dataset)}")
 
 composite = LrpLens.make_default_composite()
-cc = ChannelConcept()
 layer_names = [f"model.block{b}/conv2/relu" for b in [0, 3, 5]]
 print(layer_names)
 
@@ -96,9 +102,6 @@ def init_rel_fn(board_tensor):
 
             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:
@@ -131,23 +134,34 @@ def init_rel_fn(board_tensor):
 #######################################
 
 print("############ Clustering ...")
-os.makedirs(f"scripts/results/{run_name}", exist_ok=True)
+os.makedirs(f"scripts/results/{run_name}/{viz_name}", exist_ok=True)
 
 for layer_name, relevances in all_relevances.items():
-    kmeans = KMeans(n_clusters=n_clusters, init="k-means++")
-    kmeans.fit(relevances)
+    relevances = relevances[:n_samples]
+    if conv_sum_dims:
+        relevances = relevances.sum(dim=conv_sum_dims).view(
+            relevances.shape[0], -1
+        )
+    else:
+        relevances = relevances.view(relevances.shape[0], -1)
 
     # Perform t-SNE dimensionality reduction
     tsne = TSNE(n_components=2)
     latent_rel_tsne = tsne.fit_transform(relevances)
 
+    if kmeans_on_tsne:
+        relevances = latent_rel_tsne
+    kmeans = KMeans(n_clusters=n_clusters, init="k-means++")
+    kmeans.fit(relevances)
+
     # Plot the clustered data
     plt.scatter(latent_rel_tsne[:, 0], latent_rel_tsne[:, 1], c=kmeans.labels_)
     plt.title("Clustered Latent Relevances")
     plt.xlabel("Dimension 1")
     plt.ylabel("Dimension 2")
     plt.savefig(
-        f"scripts/results/{run_name}/{layer_name.replace('/','.')}_t-sne.png"
+        f"scripts/results/{run_name}/{viz_name}/"
+        f"{layer_name.replace('/','.')}_t-sne.png"
     )
     plt.close()
 
@@ -160,10 +174,25 @@ def init_rel_fn(board_tensor):
         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]
+            if cosine_sim:
+                dot_prod = relevances @ cluster_center.T
+                similarities = dot_prod / (
+                    np.linalg.norm(relevances, axis=1)
+                    * np.linalg.norm(cluster_center)
+                )
+                nearest_neighbors = np.argsort(similarities)[-8:]
+            else:
+                distances = np.linalg.norm(relevances - cluster_center, axis=1)
+                nearest_neighbors = np.argsort(distances)[:8]
 
-            doc = Document()  # create a new document
+            doc = Document(
+                geometry_options={
+                    "lmargin": "3cm",
+                    "tmargin": "0.5cm",
+                    "bmargin": "1.5cm",
+                    "rmargin": "3cm",
+                }
+            )
             doc.packages.append(Package("xskak"))
 
             # compute heatmap for each nearest neighbor
@@ -179,37 +208,83 @@ def init_rel_fn(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)
+
+                if viz_latent:
+                    latent_rel = all_relevances[layer_name][idx_sample]
+                    if not board.turn:
+                        latent_rel = latent_rel.flip(1)
+                    latent_rel = latent_rel.view(-1, 64)
+                    channel_rels = latent_rel.abs().sum(dim=1)
+                    c1, c2 = torch.topk(channel_rels, 2).indices
+                    heatmap_str_list = [
+                        create_heatmap_string(latent_rel.sum(0), abs_max=True),
+                        create_heatmap_string(latent_rel[c1], abs_max=True),
+                        create_heatmap_string(latent_rel[c2], abs_max=True),
+                    ]
+                    heatmap_caption_list = [
+                        "Total relevance",
+                        "Best channel",
+                        "Second best channel",
+                    ]
+                    add_caption = "latent"
+                else:
+                    board_tensor.requires_grad = True
+                    attr = attribution(
+                        board_tensor,
+                        [{"y": None}],
+                        composite,
+                        init_rel=init_rel_fn if best_legal else None,
+                    )
+                    input_relevances = board_tensor.grad
+                    if not board.turn:
+                        input_relevances = (
+                            input_relevances.view(112, 8, 8)
+                            .flip(1)
+                            .view(112, 64)
+                        )
+                    input_relevances = input_relevances.view(112, 64)
+                    heatmap_str_list = [
+                        create_heatmap_string(
+                            input_relevances.sum(dim=0), abs_max=True
+                        ),
+                        create_heatmap_string(
+                            input_relevances[:13].sum(dim=0), abs_max=True
+                        ),
+                        create_heatmap_string(
+                            input_relevances[104:].sum(dim=0), abs_max=True
+                        ),
+                    ]
+                    heatmap_caption_list = [
+                        "Total relevance",
+                        "Current config relevance",
+                        "Meta relevance",
+                    ]
+                    h0 = input_relevances[:13].abs().sum()
+                    hist = input_relevances[13:104].abs().sum()
+                    meta = input_relevances[104:].abs().sum()
+                    total = (h0 + hist + meta) / 100
+                    add_caption = (
+                        f"{h0/total:.0f}%|{hist/total:.0f}%|{meta/total:.0f}%"
+                    )
 
                 doc = add_plot(
                     doc,
                     board.fen(),
-                    heatmap_str,
+                    heatmap_str_list,
                     current_piece_pos=uci_move[:2],
                     next_move=uci_move[2:4],
+                    caption=f"Sample {idx_sample} - {add_caption}",
+                    heatmap_caption_list=heatmap_caption_list,
                 )
 
             # Generate pdf
             doc.generate_pdf(
                 f"scripts/results/{run_name}"
-                f"/{layer_name.replace('/','.')}_cluster_{idx_cluster}",
+                f"/{viz_name}/{layer_name.replace('/','.')}"
+                f"_cluster_{idx_cluster}",
                 clean_tex=True,
             )

scripts/create_figure.py

@@ -3,54 +3,70 @@
 
 
 import chess
-from pylatex import Figure, MiniPage, NoEscape
+from pylatex import Figure, NoEscape, SubFigure
 
 
 def add_plot(
     doc,
     label,
-    heatmap_str,
+    heatmap_str_list,
     current_piece_pos=None,
     next_move=None,
     caption=None,
+    heatmap_caption_list=None,
 ):
     # Put some data inside the Figure environment
     with doc.create(Figure()) as fig:
+        doc.append(NoEscape(r"\centering"))
         if caption is not None:
             fig.add_caption(caption)
         verbatim = NoEscape(
             r"\storechessboardstyle{8x8}{maxfield=h8,showmover=true}"
         )
         doc.append(verbatim)
 
-        with doc.create(MiniPage(width=r"0.45\textwidth")):
+        with doc.create(
+            SubFigure(
+                width=NoEscape(r"0.45\textwidth"),
+            )
+        ) as subfig:
+            subfig.add_caption("Board")
+            doc.append(NoEscape(r"\chessboard[style=8x8,"))
             if current_piece_pos is not None:
                 markmove = current_piece_pos + "-" + next_move
                 markfields = (
                     "{{" + current_piece_pos + "},{" + next_move + "}}"
                 )
                 chessboard_fen = NoEscape(
-                    rf"\chessboard[style=8x8,setfen={label},showmover=true,"
+                    rf"setfen={label},showmover=true,"
                     rf"color=green,pgfstyle=straightmove,markmove={markmove},"
-                    rf"pgfstyle=border,color=red,markfields={markfields},] "
+                    rf"pgfstyle=border,color=red,markfields={markfields},]"
                 )
             else:
                 chessboard_fen = NoEscape(
                     rf"\chessboard[style=8x8,setfen={label},"
-                    "showmover=true,pgfstyle=straightmove,color=green,] "
+                    "showmover=true,pgfstyle=straightmove,color=green,]"
                 )
             doc.append(chessboard_fen)
-        doc.append(NoEscape("\hfill"))  # noqa
-        with doc.create(MiniPage(width=r"0.45\textwidth")):
-            heatmap_begin = NoEscape(r"\chessboard[style=8x8,showmover=false,")
-            doc.append(heatmap_begin)
+        for i, heatmap_str in enumerate(heatmap_str_list):
+            doc.append(NoEscape(r"\hfill"))
+            with doc.create(
+                SubFigure(width=NoEscape(r"0.45\textwidth"))
+            ) as subfig:
+                subfig.add_caption(heatmap_caption_list[i])
+                heatmap_begin = NoEscape(
+                    r"\chessboard[style=8x8,showmover=false,"
+                )
+                doc.append(heatmap_begin)
 
-            heatmap_end = NoEscape(heatmap_str) + NoEscape(r"]")
-            doc.append(heatmap_end)
+                heatmap_end = NoEscape(heatmap_str) + NoEscape(r"]")
+                doc.append(heatmap_end)
     return doc
 
 
-def create_heatmap_string(heatmap):
+def create_heatmap_string(heatmap, abs_max=True):
+    if abs_max:
+        heatmap = heatmap / heatmap.abs().max()
     heatmap_str = ""
     for idx, name in enumerate(chess.SQUARE_NAMES):
         colorcode = heatmap[idx]