File size: 2,345 Bytes
a529eec | 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 | 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 # expected shape (M,), dtype long, on same device
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:
# Skip partial batches when iterating across epoch boundaries
if len(batch) < self.batch_size:
continue
yield batch
steps_yielded += 1
if steps_yielded >= n_steps:
return
|