| import torch |
| from typing import Iterator, Optional |
|
|
|
|
| class TensorBuffer: |
| """Shuffled batch iterator over a tensor of activations. |
| |
| When ``subset_indices`` is provided, iterates only over ``data[subset_indices]`` |
| without materialising a copy — batches are gathered per-step. This keeps |
| memory at O(batch_size) per step instead of O(N) per epoch, which matters |
| for large activation tensors at high expansion ratios. |
| """ |
|
|
| def __init__( |
| self, |
| data: torch.Tensor, |
| batch_size: int = 4096, |
| subset_indices: Optional[torch.Tensor] = None, |
| ) -> None: |
| self.data = data |
| self.batch_size = batch_size |
| self.subset_indices = subset_indices |
|
|
| def __iter__(self) -> Iterator[torch.Tensor]: |
| """Shuffle (subset of) data and yield batches for one epoch. |
| |
| Emits batch-sized gathers rather than a full shuffled copy, so peak |
| memory stays near O(batch_size * D) regardless of N. |
| """ |
| n = self.data.shape[0] if self.subset_indices is None else self.subset_indices.shape[0] |
| perm = torch.randperm(n, device="cpu") |
| for start in range(0, n, self.batch_size): |
| batch_perm = perm[start : start + self.batch_size] |
| if self.subset_indices is None: |
| idx = batch_perm.to(self.data.device) |
| else: |
| idx = self.subset_indices.index_select(0, batch_perm.to(self.subset_indices.device)) |
| yield self.data.index_select(0, idx) |
|
|
| def iter_steps(self, n_steps: int) -> Iterator[torch.Tensor]: |
| """Yield exactly n_steps full batches, wrapping around with reshuffle each epoch.""" |
| n = self.data.shape[0] if self.subset_indices is None else self.subset_indices.shape[0] |
| if self.batch_size > n: |
| raise ValueError( |
| f"batch_size ({self.batch_size}) > data size ({n})" |
| ) |
| steps_yielded = 0 |
| while steps_yielded < n_steps: |
| for batch in self: |
| |
| if len(batch) < self.batch_size: |
| continue |
| yield batch |
| steps_yielded += 1 |
| if steps_yielded >= n_steps: |
| return |
|
|
