omniloader.sampling.bucketing

Length-bucketed batching to minimize padding waste.

LengthBucketBatchSampler wraps an index sampler (e.g. OmniSampler) and a per-sample length array, and yields batches of indices whose sequences have similar length. Used as DataLoader(batch_sampler=...) together with DynamicCollator, it keeps each batch’s padded length close to its true length.

The classic pooling scheme is used: the sampler’s index stream is cut into pools of batch_size * bucket_multiplier; each pool is sorted by length and split into batches; the batch order is then shuffled so length still varies across steps (only within-batch length is homogenised).

This only reduces padding waste in the native-length regime (OmniLoader(pad_features=False) + DynamicCollator). When every sample is padded/cropped to a fixed time_dim the batches are already equal length, so grouping by length changes no shapes and bucketing is a no-op.

Classes

SizedSampler	An index sampler that reports its length (as PyTorch samplers do).
LengthBucketBatchSampler	Yield batches of indices grouped by sequence length.

Module Contents

class omniloader.sampling.bucketing.SizedSampler

Bases: Protocol

An index sampler that reports its length (as PyTorch samplers do).

__iter__() → collections.abc.Iterator[int]

Iterate the sampler’s indices.

__len__() → int

Number of indices produced per epoch.

class omniloader.sampling.bucketing.LengthBucketBatchSampler(sampler: SizedSampler, lengths: collections.abc.Sequence[int], batch_size: int, bucket_multiplier: int = 100, drop_last: bool = False, shuffle: bool = True, seed: int = 0)

Bases: torch.utils.data.Sampler[list[int]]

Yield batches of indices grouped by sequence length.

Only meaningful in the native-length regime (pad_features=False + DynamicCollator); under a fixed time_dim all batches are already equal length and this changes nothing.

Parameters:

• sampler – An index sampler yielding global indices (its length defines the epoch size). If it exposes an epoch attribute, it seeds the batch shuffle for reproducibility across epochs.

• lengths – Per-index sequence lengths, indexable by the sampler’s indices (e.g. from sequence_lengths()).

• batch_size – Number of samples per batch.

• bucket_multiplier – Pool size in batches; larger pools bucket more tightly but reduce cross-length shuffling.

• drop_last – Drop a trailing batch smaller than batch_size.

• shuffle – Shuffle the order of the produced batches each epoch.

• seed – Base seed for the batch-order shuffle.

Raises:

ValueError – If batch_size or bucket_multiplier is not positive.

__iter__() → collections.abc.Iterator[list[int]]

Iterate length-bucketed batches for the current epoch.

__len__() → int

Number of batches produced per epoch.