quantlab.model.torch_data.IndexBatches#

class quantlab.model.torch_data.IndexBatches(size, batch_size, *, shuffle=False, drop_last=False, replacement=False, num_samples=None, generator=None)[source]#

Bases: Sampler

Batches of dataset positions, each an int64 tensor, for a dataset that gathers a batch at once.

Each pass draws a fresh order: positions in order, shuffled (torch.randperm), or num_samples drawn with replacement (torch.randint), from generator; consecutive batch_size positions form a batch, and drop_last drops a last batch shorter than batch_size.

Parameters:

sizeint: Number of items in the dataset.
batch_sizeint: Positions per batch; at least 1.
shufflebool, default False: Draw a new random order every pass.
drop_lastbool, default False: Drop a last incomplete batch.
replacementbool, default False: Draw num_samples positions with replacement (implies shuffling).
num_samplesint, optional: Positions drawn per pass with replacement; size when unset.
generatortorch.Generator, optional: Source of the random order.

Raises:

ValueError: If batch_size is below 1.

Examples

>>> [b.tolist() for b in IndexBatches(5, 2)]
[[0, 1], [2, 3], [4]]
>>> len(IndexBatches(5, 2, drop_last=True))
2