Add CachedActivationStore

Author: samuelstevens

analysis.py

@@ -5,12 +5,12 @@
 
 import beartype
 import torch
-import tqdm
 import tyro
 from jaxtyping import Float, Int, jaxtyped
 from torch import Tensor
 
 import saev
+from saev import helpers
 
 # Fix pickle renaming errors.
 sys.modules["sae_training"] = saev
@@ -39,7 +39,7 @@ def batched_idx(
 
 @jaxtyped(typechecker=beartype.beartype)
 def get_vit_acts(
-    acts_store: saev.ActivationsStore, n: int
+    acts_store: saev.CachedActivationsStore, n: int
 ) -> tuple[Float[Tensor, "n d_model"], Int[Tensor, " n"]]:
     """
     Args:
@@ -54,7 +54,6 @@ def get_vit_acts(
         batches.append(batch)
         indices.append(i)
         n_seen += len(batch)
-        logger.info("Got batch of size %d (%d total).", len(batch), n_seen)
 
     batches = torch.cat(batches, dim=0)
     indices = torch.cat(indices, dim=0)
@@ -94,7 +93,7 @@ def get_new_topk(
 @torch.inference_mode()
 def get_feature_data(
     sae: saev.SparseAutoencoder,
-    acts_store: saev.ActivationsStore,
+    acts_store: saev.CachedActivationsStore,
     *,
     n_images: int = 32_768,
     k_top_images: int = 10,
@@ -112,14 +111,14 @@ def get_feature_data(
     torch.cuda.empty_cache()
     sae.eval()
 
-    if n_images > len(acts_store.dataset):
+    if n_images > len(acts_store):
         logger.warning(
             "The dataset '%s' only has %d images, but you requested %d images.",
             sae.cfg.dataset_path,
-            len(acts_store.dataset),
+            len(acts_store),
             n_images,
         )
-        n_images = len(acts_store.dataset)
+        n_images = len(acts_store)
 
     top_values = torch.zeros((sae.cfg.d_sae, k_top_images)).to(sae.cfg.device)
     top_indices = torch.zeros((sae.cfg.d_sae, k_top_images), dtype=torch.int)
@@ -128,13 +127,18 @@ def get_feature_data(
     sae_sparsity = torch.zeros((sae.cfg.d_sae,)).to(sae.cfg.device)
     sae_mean_acts = torch.zeros((sae.cfg.d_sae,)).to(sae.cfg.device)
 
-    n_seen = 0
+    dataloader = torch.utils.data.DataLoader(
+        acts_store,
+        batch_size=images_per_it,
+        shuffle=True,
+        num_workers=sae.cfg.n_workers,
+        drop_last=True,
+    )
 
-    while n_seen < n_images:
+    for batch in helpers.progress(dataloader):
         torch.cuda.empty_cache()
+        vit_acts, indices = batch
 
-        # tensor of size [batch, d_resid]
-        vit_acts, indices = get_vit_acts(acts_store, images_per_it)
         # tensor of size [feature_idx, batch]
         sae_acts = get_sae_acts(vit_acts.to(sae.cfg.device), sae).transpose(0, 1)
         del vit_acts
@@ -149,24 +153,18 @@ def get_feature_data(
             top_values, top_indices, values, indices, k_top_images
         )
 
-        n_seen += images_per_it
-        logger.info("%d/%d (%.1f%%)", n_seen, n_images, n_seen / n_images * 100)
-
     sae_mean_acts /= sae_sparsity
-    sae_sparsity /= n_images
+    sae_sparsity /= len(acts_store)
 
     # Check if the directory exists
     if not os.path.exists(directory):
         # Create the directory if it does not exist
         os.makedirs(directory)
 
-    # compute the label tensor
-    top_image_label_indices = torch.tensor([
-        acts_store.dataset[int(index)]["label"]
-        for index in tqdm.tqdm(top_indices.flatten(), desc="Getting labels")
-    ])
-    # Reshape to original dimensions
-    top_image_label_indices = top_image_label_indices.view(top_indices.shape)
+    # Compute the label tensor
+    top_image_label_indices = acts_store.labels[top_indices.view(-1).cpu()].view(
+        top_indices.shape
+    )
     torch.save(top_indices, f"{directory}/max_activating_image_indices.pt")
     torch.save(top_values, f"{directory}/max_activating_image_values.pt")
     torch.save(

generate_app_data.py

@@ -20,6 +20,32 @@ def safe_load(path: str) -> object:
     return torch.load(path, map_location="cpu", weights_only=True)
 
 
+@beartype.beartype
+def make_img_grid(imgs: list):
+    # Resize to 224x224
+    img_width, img_height = 224, 224
+    imgs = [img.resize((img_width, img_height)).convert("RGB") for img in imgs]
+
+    # Create an image grid
+    grid_size = 4
+    border_size = 2  # White border thickness
+
+    # Create a new image with white background
+    grid_width = grid_size * img_width + (grid_size - 1) * border_size
+    grid_height = grid_size * img_height + (grid_size - 1) * border_size
+    img_grid = Image.new("RGB", (grid_width, grid_height), "white")
+
+    # Paste images in the grid
+    x_offset, y_offset = 0, 0
+    for i, img in enumerate(imgs):
+        img_grid.paste(img, (x_offset, y_offset))
+        x_offset += img_width + border_size
+        if (i + 1) % grid_size == 0:
+            x_offset = 0
+            y_offset += img_height + border_size
+    return img_grid
+
+
 @beartype.beartype
 def main(ckpt_path: str, in_dir: str = "dashboard", out_dir: str = "web_app"):
     """
@@ -85,31 +111,6 @@ def main(ckpt_path: str, in_dir: str = "dashboard", out_dir: str = "web_app"):
     indices = torch.tensor([i for i in range(n_neurons)])
     indices = list(indices[mask])
 
-    @beartype.beartype
-    def make_img_grid(imgs: list):
-        # Resize to 224x224
-        img_width, img_height = 224, 224
-        imgs = [img.resize((img_width, img_height)).convert("RGB") for img in imgs]
-
-        # Create an image grid
-        grid_size = 4
-        border_size = 2  # White border thickness
-
-        # Create a new image with white background
-        grid_width = grid_size * img_width + (grid_size - 1) * border_size
-        grid_height = grid_size * img_height + (grid_size - 1) * border_size
-        img_grid = Image.new("RGB", (grid_width, grid_height), "white")
-
-        # Paste images in the grid
-        x_offset, y_offset = 0, 0
-        for i, img in enumerate(imgs):
-            img_grid.paste(img, (x_offset, y_offset))
-            x_offset += img_width + border_size
-            if (i + 1) % grid_size == 0:
-                x_offset = 0
-                y_offset += img_height + border_size
-        return img_grid
-
     os.makedirs(f"{out_dir}/neurons", exist_ok=True)
     torch.save(entropies, f"{out_dir}/neurons/entropy.pt")
     for i in tqdm.tqdm(indices, desc="saving highest activating grids"):

logbook.md

@@ -89,8 +89,8 @@ With this in mind, there are several minor changes I want to make before I do so
 
 1. Removing `transformer-lens` [done, commit [18612b7](https://github.com/samuelstevens/saev/commit/18612b75988c32ae8ab3db6656b44a442f3f7641)]
 2. Removing HookedVisionTransformer [done, commit [c7ba7c7](https://github.com/samuelstevens/saev/commit/c7ba7c72c76472fd8cf2e7b2dc668d03a15b803d)]
-3. OpenCLIP instead of huggingface `transformers` [done, testing]
-4. Pre-computing ViT activations
+3. OpenCLIP instead of huggingface `transformers` [done, commit [d362f64](https://github.com/samuelstevens/saev/commit/d362f64437b3599f56bb698136712d7590ee897b)]
+4. Pre-computing ViT activations [done, commit [ee79f5b](https://github.com/samuelstevens/saev/commit/ee79f5b84186e655b2e5d485e972fe69bb73dd65)]
 
 I'm going to do each of these independently using a set of runs as references.
 
@@ -111,3 +111,23 @@ Only after that will I use the new class in training.
 Working with the analysis script is a shorter feedback loop.
 
 # 10/23/2024
+
+OpenCLIP instead of transformers works (training, analysis, generate).
+So now I am pre-computing activations.
+I'm waiting on the activations to be saved (~3 hours).
+
+CachedActivationsStore produced some duplicates in the analysis step.
+Why is that?
+
+For example, neuron 78 has the same image for image 6 and 7 (1-indexed, images 5 and 6 if zero-indexed).
+
+Fixed it.
+We no longer randomly sample batches; instead, we use a dataloader and `__getitem__`.
+
+With training, however, the metrics no longer match the reference metrics.
+Why is that?
+We can find out by comparing to the original activations store.
+Likely, we will need to build a custom data order using `np.random.default_rng(seed=cfg.seed)`.
+
+My strategy for calculating the mean activations only used 15 examples instead of 15 x 1024.
+With 15 x 1024 examples, the b_dec is better initialized and it works exactly like before.

saev/__init__.py

@@ -1,5 +1,6 @@
 from .modeling import (
     ActivationsStore,
+    CachedActivationsStore,
     Config,
     RecordedVit,
     Session,
@@ -9,6 +10,7 @@
 
 __all__ = [
     "ActivationsStore",
+    "CachedActivationsStore",
     "Config",
     "RecordedVit",
     "SparseAutoencoder",

saev/helpers.py

@@ -0,0 +1,53 @@
+import collections.abc
+import logging
+import time
+
+import beartype
+
+
+@beartype.beartype
+class progress:
+    def __init__(self, it, *, every: int = 10, desc: str = "progress", total: int = 0):
+        """
+        Wraps an iterable with a logger like tqdm but doesn't use any control codes to manipulate a progress bar, which doesn't work well when your output is redirected to a file. Instead, simple logging statements are used, but it includes quality-of-life features like iteration speed and predicted time to finish.
+
+        Args:
+            it: Iterable to wrap.
+            every: How many iterations between logging progress.
+            desc: What to name the logger.
+            total: If non-zero, how long the iterable is.
+        """
+        self.it = it
+        self.every = every
+        self.logger = logging.getLogger(desc)
+        self.total = total
+
+    def __iter__(self):
+        start = time.time()
+        for i, obj in enumerate(self.it):
+            yield obj
+
+            if (i + 1) % self.every == 0:
+                now = time.time()
+                duration_s = now - start
+                per_min = (i + 1) / (duration_s / 60)
+
+                if isinstance(self.it, collections.abc.Sized):
+                    pred_min = (len(self) - (i + 1)) / per_min
+                    self.logger.info(
+                        "%d/%d (%.1f%%) | %.1f it/m (expected finish in %.1fm)",
+                        i + 1,
+                        len(self),
+                        (i + 1) / len(self) * 100,
+                        per_min,
+                        pred_min,
+                    )
+                else:
+                    self.logger.info("%d/? | %.1f it/m", i + 1, per_min)
+
+    def __len__(self) -> int:
+        if self.total > 0:
+            return self.total
+
+        # Will throw exception.
+        return len(self.it)