DataLoader
Iterate a dataset in batches, optionally shuffled and augmented.
Single-process and synchronous: batches are assembled on the calling thread. Prefetching and worker processes are tracked as future work, not implemented.
transform= takes a callable f(rng, arrays) -> arrays, applied to each batch before it is moved to the device. Two things about that signature are deliberate:
- It receives the whole batch, because
Dataset.__getitem__is batched too. A flip is one vectorised operation over 64 images rather than 64 Python calls. - The generator is passed in, never reached for. Augmentation that draws from NumPy's global state is irreproducible, and this project has already paid for exactly that once —
nn.manual_seedexists because layers called an unseededdefault_rng()and a divergence could not be re-observed. A transform would have to import NumPy and ignore its argument to become non-deterministic.
Shuffling is seeded, so a run replays. The transform draws from a separate stream: sharing one with the shuffle would make the same seed give different augmentation depending on whether shuffling was on.
The final batch is smaller when the dataset size is not a multiple of the batch size — it is not dropped unless drop_last says so, so every sample is seen exactly once per epoch.
Augmentation is real host work: it takes CIFAR-100's batch-production time from 1.7% of a step to 10.6%, measured. That is also the number that decides whether prefetching is worth building — and prefetching turns out not to be a DataLoader change at all, because the bindings hold the GIL through the GPU wait and a producer thread would get only a fifth of the window it needs.
Construction¶
__init__¶
Iteration¶
__iter__¶
__len__¶
Number of batches one pass yields.
Other members¶
__repr__¶
set_epoch¶
Pin the shuffle to a specific epoch.
Iteration advances this on its own; setting it explicitly is what lets a run resume mid-training and see the order it would have seen.
See also ArrayDataset, Compose, rand