Spaces:
Sleeping
Sleeping
File size: 5,229 Bytes
bda104d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 | """
SaveActivationsCallback — saves validation embeddings, preactivations,
activations, and reconstructions to per-tensor .pt files at the end of each
validation epoch.
Reads embeddings and activations stored on the LightningModule by _eval_step,
then re-runs the SAE forward pass to capture preactivations and reconstructions
(which are not stored on the module to save memory during normal training).
"""
from __future__ import annotations
import logging
import os
from typing import Dict, Optional
import torch
from lightning.pytorch import Trainer
from lightning.pytorch.core import LightningModule
from steerable_retrieval.callbacks.utils import BaseCallback
from steerable_retrieval.callbacks.energy import _get_dataset_name, gather_tensor_if_distributed, _cat_or_none
log = logging.getLogger(__name__)
class SaveActivationsCallback(BaseCallback):
"""
At the end of every validation epoch, saves separate .pt files containing
embeddings, preactivations, activations, and reconstructions for every
validation dataloader / modality.
File structure::
<save_dir>/activations/<dataset_name>/<modality>/<tensor_name>/tensors.pt
Example::
activations/dataset_name/audio/embeddings/tensors.pt
activations/dataset_name/audio/preactivations/tensors.pt
activations/dataset_name/audio/activations/tensors.pt
activations/dataset_name/audio/reconstructions/tensors.pt
"""
ROOT_DIRNAME = "activations"
def __init__(
self,
save_dir: str,
every_n_steps: int = None,
every_n_epochs: int = 1,
):
super().__init__(every_n_steps=every_n_steps, every_n_epochs=every_n_epochs)
self.save_dir = save_dir
def on_validation_epoch_end(self, trainer: Trainer, pl_module: LightningModule):
if not (self._check_step(trainer, pl_module) or self._check_epoch(trainer, pl_module)):
return
self._save(trainer, pl_module)
@torch.no_grad()
def _save(self, trainer: Trainer, pl_module: LightningModule):
"""Build tensors from module-stored embeddings/activations and save to disk."""
device = pl_module.device
all_acts = getattr(pl_module, "val_activations", {})
all_embs = getattr(pl_module, "val_embeddings", {})
for dataloader_idx in all_acts:
dataset_name = _get_dataset_name(trainer, dataloader_idx, "val")
acts_dl = all_acts.get(dataloader_idx, {})
embs_dl = all_embs.get(dataloader_idx, {})
for modality in ("audio", "text"):
act_chunks = acts_dl.get(modality, [])
emb_chunks = embs_dl.get(modality, [])
if not act_chunks:
continue
# Concatenate and gather
activations = _cat_or_none(act_chunks, device=device)
embeddings = _cat_or_none(emb_chunks, device=device) if emb_chunks else None
activations = gather_tensor_if_distributed(activations, trainer)
if embeddings is not None:
embeddings = gather_tensor_if_distributed(embeddings, trainer)
# Re-run SAE forward on embeddings to get preactivations & reconstructions
preactivations = None
reconstructions = None
if embeddings is not None:
# Process in chunks to avoid OOM
chunk_size = 512
pre_chunks, rec_chunks = [], []
for i in range(0, embeddings.size(0), chunk_size):
emb_chunk = embeddings[i : i + chunk_size].to(device)
_, z, xhat, pre = pl_module(emb_chunk)
pre_chunks.append(pre.cpu() if pre is not None else torch.zeros_like(z).cpu())
rec_chunks.append(xhat.cpu())
preactivations = torch.cat(pre_chunks, dim=0)
reconstructions = torch.cat(rec_chunks, dim=0)
tensors = {
"embeddings": embeddings.cpu() if embeddings is not None else None,
"preactivations": preactivations,
"activations": activations.cpu(),
"reconstructions": reconstructions,
}
if trainer.is_global_zero:
self._save_tensors(dataset_name=dataset_name, modality=modality, tensors=tensors)
# Only rank-0 writes to disk
if trainer.is_global_zero:
log.info(f"[SaveActivationsCallback] Saved validation activations to {self.save_dir}")
def _save_tensors(
self,
dataset_name: str,
modality: str,
tensors: Dict[str, Optional[torch.Tensor]],
) -> None:
base_dir = os.path.join(self.save_dir, self.ROOT_DIRNAME, dataset_name, modality)
for tensor_name, tensor_value in tensors.items():
if tensor_value is None:
continue
tensor_dir = os.path.join(base_dir, tensor_name)
os.makedirs(tensor_dir, exist_ok=True)
tensor_path = os.path.join(tensor_dir, "tensors.pt")
torch.save(tensor_value, tensor_path)
|