create heatmap-visualizations for sweeps

elk/training/sweep.py

@@ -3,6 +3,8 @@
 import numpy as np
 import torch
 
+from elk.utils.visualize import render_model_results
+
 from ..evaluation import Eval
 from ..extraction import Extract
 from ..files import elk_reporter_dir, memorably_named_dir
@@ -29,6 +31,9 @@ class Sweep:
     If negative, no hyperparameter sweeps will be performed. Only valid for Eigen."""
     skip_transfer_eval: bool = False
     """Whether to perform transfer eval on every pair of datasets."""
+
+    visualize: bool = False
+    """Whether to generate visualizions of the results of the sweep."""
 
     name: str | None = None
 
@@ -163,3 +168,7 @@ def execute(self):
                                     skip_supervised=run.supervised == "none",
                                 )
                                 eval.execute(highlight_color="green")
+
+        if self.visualize:
+            for i, model in enumerate(self.models):
+                render_model_results(sweep_dir / model, sweep_dir / "visualizations")

elk/utils/visualize.py

@@ -0,0 +1,72 @@
+import os
+from pathlib import Path
+
+import pandas as pd
+import plotly.express as px
+
+
+def generate_heatmap(data, model_name, layer, viz_folder_path):
+
+    pivot = pd.pivot_table(
+        data, values="auroc_estimate", index="eval_dataset", columns="train_dataset"
+    )
+
+    fig = px.imshow(pivot, color_continuous_scale="Viridis", text_auto=True)
+
+    fig.update_layout(
+        xaxis_title="Train Dataset",
+        yaxis_title="Transfer Dataset",
+        title=f"AUROC Score Heatmap: {model_name} | Layer {layer}",
+    )
+
+    model_viz_folder_path = os.path.join(viz_folder_path, model_name)
+    if not os.path.exists(model_viz_folder_path):
+        os.mkdir(model_viz_folder_path)
+    fig.write_image(os.path.join(model_viz_folder_path, f"{layer}.png"))
+
+
+def filter_df(df, layer):
+    df = df[df["ensembling"] == "full"]  # filter df to only include full ensembling
+    df = df[df["layer"] == layer]  # filter df to only include target layer
+    return df
+
+
+def render_model_results(model_path, visualizations_path):
+    if not visualizations_path.exists():
+        visualizations_path.mkdir()
+
+    df = None
+
+    def get_layer_min_max(model_path: Path):
+        dir = model_path.iterdir().__next__()
+        file = os.path.join(dir, "eval.csv")
+        raw_eval_df = pd.read_csv(file)