DataLoader
DataLoader(dataset: Dataset, batch_size: int | None = 1, shuffle: bool | None = None, sampler: Sampler | None = None, batch_sampler: Sampler | None = None, num_workers: int = 0, collate_fn: Callable[..., object] | None = None, drop_last: bool = False, timeout: float = 0.0, worker_init_fn: Callable[..., object] | None = None, multiprocessing_context: object = None, generator: Generator | int | None = None, prefetch_factor: int | None = None, persistent_workers: bool = False, pin_memory: bool = False)Combine a dataset with a sampler to provide iteration over mini-batches.
Wraps a Dataset to provide batching, optional shuffling,
parallel data loading via worker processes, and customisable
collation. Iteration yields one collated batch per step until the
underlying sampler is exhausted.
Parameters
datasetDatasetDataset) or iterable-style (IterableDataset).batch_sizeint or None= 1batch_sampler is
provided. None turns automatic batching off: each step
yields one sample, passed through collate_fn
(default_convert by default).shufflebool= NoneTrue, the default sampler is RandomSampler;
otherwise SequentialSampler. Mutually exclusive with
sampler.shuffle.batch_size / shuffle / sampler /
drop_last.num_workersint= 00 runs
single-process in the main thread; > 0 spawns a worker pool.
With an IterableDataset every worker iterates its own
copy of the dataset — shard it with get_worker_info, or
each sample arrives num_workers times.collate_fncallable= Nonedefault_collate, or default_convert when
batch_size=None).drop_lastbool= FalseTrue, drop the trailing batch when the dataset length is
not divisible by batch_size.timeoutfloat= 0.0RuntimeError. 0 waits indefinitely — for a live worker:
a worker that dies raises either way.worker_init_fncallable= Noneworker_init_fn(worker_id) in each worker process,
after the worker's random generators are seeded and
get_worker_info is available.multiprocessing_contextstr or context= None"spawn".RandomSampler's order and the workers' base seed.
None (default) uses the global generator, so
lucid.manual_seed reproduces both.prefetch_factorint= None2 when
num_workers > 0). Higher values trade memory for throughput.persistent_workersbool= Falsenum_workers > 0.pin_memorybool= FalseNotes
Worker processes communicate via multiprocessing.Queue: each
worker owns one index queue, all workers share a single result
queue, and the main process reorders results back into sampler
order before yielding. Sequence numbers ensure deterministic
delivery regardless of completion order across workers.
Randomness in workers. Each iterator draws one base_seed
from generator (or the global generator), and worker i seeds
Lucid's generator, Python's random and — when installed —
NumPy's legacy global RNG with base_seed + i before
worker_init_fn runs. Random augmentations therefore differ
between workers and between epochs, and the whole run reproduces
under lucid.manual_seed. A persistent pool is seeded once
and its streams carry on from epoch to epoch.
An exception raised in a worker is re-raised in the main process, as the same exception type where possible, with the worker's traceback in its message.
Examples
>>> import lucid
>>> from lucid.utils.data import DataLoader, TensorDataset
>>> dataset = TensorDataset(lucid.randn(100, 8), lucid.randint(0, 10, (100,)))
>>> dl = DataLoader(dataset, batch_size=32, shuffle=True) # num_workers=4 to parallelise
>>> for batch in dl:
... x, y = batch
>>> x.shape, y.shape # the last, short batch: 100 = 3 * 32 + 4
((4, 8), (4,))
One sample per step, without automatic batching:
>>> x, y = next(iter(DataLoader(dataset, batch_size=None)))
>>> x.shape, y.shape
((8,), ())Used by 1
Constructors
1__init__
→None__init__(dataset: Dataset, batch_size: int | None = 1, shuffle: bool | None = None, sampler: Sampler | None = None, batch_sampler: Sampler | None = None, num_workers: int = 0, collate_fn: Callable[..., object] | None = None, drop_last: bool = False, timeout: float = 0.0, worker_init_fn: Callable[..., object] | None = None, multiprocessing_context: object = None, generator: Generator | int | None = None, prefetch_factor: int | None = None, persistent_workers: bool = False, pin_memory: bool = False)Configure a DataLoader; see the class docstring for parameter
semantics.
Raises
ValueErrorbatch_size or prefetch_factor that is not a positive
integer (or None); a num_workers that is not an integer
>= 0; a negative or NaN timeout; prefetch_factor,
multiprocessing_context or persistent_workers given
with num_workers=0; or an unknown start-method name.TypeErrortimeout is not a number, collate_fn or
worker_init_fn is not callable, multiprocessing_context
is neither a name nor a context, or generator is neither a
lucid.Generator, a seed, nor None (a bool is
not a seed).Notes
sampler / batch_sampler / shuffle / batch_size /
drop_last interact: passing batch_sampler precludes the
other four; passing sampler precludes shuffle. When no
sampler is supplied, a SequentialSampler (shuffle=False)
or RandomSampler (shuffle=True) is constructed
automatically. persistent_workers requires num_workers > 0;
batch_size=None precludes drop_last.
Dunder methods
2Return an iterator over one full pass of collated mini-batches.
Dispatches to either the single-process iterator (num_workers == 0) or the multi-process iterator. When persistent_workers is
enabled the multi-process worker pool survives between epochs and
the same iterator is reset for each pass; otherwise workers are
spawned per pass and shut down when it ends or the iterator is
dropped.
Returns
IteratorYields the output of collate_fn applied to each sampled
batch of dataset items.
Return the number of steps per epoch (len of the index sampler).
An iterable-style dataset has no length to divide, so neither has
the loader over it. TypeError rather than a guess: a wrong
len silently truncates a progress bar, a learning-rate
schedule, or an epoch.