| """Frame-budget dynamic batching: deterministic, bucketed, resumable. |
| |
| An *epoch* visits every item exactly once. For epoch ``e`` the sampler |
| |
| 1. draws a permutation of all items from ``default_rng([seed, stream, e])``; |
| 2. cuts it into pools of ``pool_size`` items and sorts each pool by length (stable, so ties |
| keep the shuffled order) -- items inside a batch therefore have similar length; |
| 3. greedily packs each sorted pool into batches whose padded size |
| ``max_len * batch_size`` stays ``<= frame_budget`` (and ``<= max_items`` items); |
| 4. shuffles the order of the batches. |
| |
| The result depends only on ``(lengths, frame_budget, pool_size, max_items, seed, stream, |
| epoch)``, so a run can be resumed from a :class:`SamplerPosition` ``(epoch, batch_index)``. |
| Nothing is ever dropped or split: an item longer than the budget raises |
| :class:`BudgetError` at construction. |
| |
| ``stream`` separates independent uses (train = 0, eval = 1) that share one ``seed``. |
| """ |
|
|
| from __future__ import annotations |
|
|
| from collections.abc import Sequence |
| from dataclasses import dataclass |
|
|
| import numpy as np |
|
|
| __all__ = ["BatchStream", "BudgetError", "FrameBudgetSampler", "SamplerPosition"] |
|
|
|
|
| class BudgetError(ValueError): |
| """An item does not fit into the frame budget (it would have to be split or dropped).""" |
|
|
|
|
| @dataclass(frozen=True) |
| class SamplerPosition: |
| """Where the stream is: the *next* batch to be produced is ``batch_index`` of ``epoch``.""" |
|
|
| epoch: int = 0 |
| batch_index: int = 0 |
|
|
|
|
| class FrameBudgetSampler: |
| """Deterministic epoch -> list of batches (each batch: item indices). |
| |
| Args: |
| lengths: ``num_frames`` of every item, in dataset order (all >= 1). |
| frame_budget: maximum padded frames per batch (``max_len * batch_size``). |
| seed: non-negative base seed. |
| pool_size: items per sorting pool (bigger = more similar lengths, less randomness). |
| max_items: optional cap on the batch size. |
| stream: independent random stream id (non-negative). |
| |
| Raises: |
| ValueError: invalid arguments or an empty dataset. |
| BudgetError: some item is longer than ``frame_budget``. |
| """ |
|
|
| def __init__( |
| self, |
| lengths: Sequence[int], |
| *, |
| frame_budget: int, |
| seed: int, |
| pool_size: int, |
| max_items: int | None = None, |
| stream: int = 0, |
| ) -> None: |
| if len(lengths) == 0: |
| raise ValueError("cannot sample from an empty dataset") |
| if min(lengths) < 1: |
| raise ValueError("every length must be >= 1") |
| if seed < 0 or stream < 0: |
| raise ValueError("seed and stream must be >= 0") |
| if frame_budget < 1 or pool_size < 1 or (max_items is not None and max_items < 1): |
| raise ValueError("frame_budget, pool_size and max_items must be >= 1") |
| too_long = [(i, n) for i, n in enumerate(lengths) if n > frame_budget] |
| if too_long: |
| index, frames = too_long[0] |
| raise BudgetError( |
| f"{len(too_long)} item(s) exceed frame_budget={frame_budget}; " |
| f"first: item {index} with {frames} frames" |
| ) |
| self._lengths = np.asarray(lengths, dtype=np.int64) |
| self.frame_budget = frame_budget |
| self.seed = seed |
| self.pool_size = pool_size |
| self.max_items = max_items |
| self.stream = stream |
| self._cache: tuple[int, list[list[int]]] | None = None |
|
|
| @property |
| def num_items(self) -> int: |
| return int(self._lengths.shape[0]) |
|
|
| def epoch_batches(self, epoch: int) -> list[list[int]]: |
| """All batches of ``epoch`` (item indices; each item appears exactly once).""" |
| if epoch < 0: |
| raise ValueError("epoch must be >= 0") |
| if self._cache is not None and self._cache[0] == epoch: |
| return self._cache[1] |
| rng = np.random.default_rng([self.seed, self.stream, epoch]) |
| order = rng.permutation(self.num_items) |
| batches: list[list[int]] = [] |
| for start in range(0, self.num_items, self.pool_size): |
| pool = order[start : start + self.pool_size] |
| pool = pool[np.argsort(self._lengths[pool], kind="stable")] |
| current: list[int] = [] |
| for index in pool.tolist(): |
| length = int(self._lengths[index]) |
| full = bool(current) and ( |
| length * (len(current) + 1) > self.frame_budget |
| or (self.max_items is not None and len(current) >= self.max_items) |
| ) |
| if full: |
| batches.append(current) |
| current = [] |
| current.append(index) |
| if current: |
| batches.append(current) |
| permutation = rng.permutation(len(batches)) |
| result = [batches[i] for i in permutation.tolist()] |
| self._cache = (epoch, result) |
| return result |
|
|
| def num_batches(self, epoch: int) -> int: |
| return len(self.epoch_batches(epoch)) |
|
|
|
|
| class BatchStream: |
| """Infinite stream of batches over epochs with a resumable position.""" |
|
|
| def __init__( |
| self, sampler: FrameBudgetSampler, position: SamplerPosition | None = None |
| ) -> None: |
| self._sampler = sampler |
| self._epoch = 0 |
| self._batch_index = 0 |
| if position is not None: |
| self.set_position(position) |
|
|
| @property |
| def position(self) -> SamplerPosition: |
| return SamplerPosition(self._epoch, self._batch_index) |
|
|
| def set_position(self, position: SamplerPosition) -> None: |
| """Jump to ``position``. |
| |
| Raises: |
| ValueError: ``batch_index`` is beyond the number of batches of that epoch. |
| """ |
| count = self._sampler.num_batches(position.epoch) |
| if not 0 <= position.batch_index <= count: |
| raise ValueError( |
| f"batch_index {position.batch_index} outside 0..{count} for epoch {position.epoch}" |
| ) |
| self._epoch, self._batch_index = position.epoch, position.batch_index |
|
|
| def next_batch(self) -> list[int]: |
| """Item indices of the next batch; rolls over to the next epoch when exhausted.""" |
| batches = self._sampler.epoch_batches(self._epoch) |
| if self._batch_index >= len(batches): |
| self._epoch += 1 |
| self._batch_index = 0 |
| batches = self._sampler.epoch_batches(self._epoch) |
| batch = batches[self._batch_index] |
| self._batch_index += 1 |
| return list(batch) |
|
|
