Tracking: DataLoader feature parity with PyTorch
- Dominant language
- Julia
- Stars
- 124
- Forks
- 23
- PR merge metrics
- No merged PRs in 30d
Description
## Tracking: `DataLoader` feature parity with PyTorch `torch.utils.data.DataLoader`
Today `MLUtils.DataLoader` exposes `batchsize, buffer, collate, parallel, num_workers, partial, rng, shuffle`. Comparing against PyTorch's `DataLoader`, the following features are missing. Filing as a tracking issue to gauge interest and scope — not all are necessarily in scope for MLUtils core.
### A. Custom samplers (`sampler` / `batch_sampler`)
PyTorch lets you replace iteration order with a `sampler` (yields indices, e.g. `WeightedRandomSampler`, `SubsetRandomSampler`) or a `batch_sampler` (yields lists of indices → full control over batch composition and variable batch sizes). MLUtils offers only `shuffle=true/false`. Since index sets are already computed on the main process, a `sampler` hook slots in cleanly. Missing:
- Weighted / subset / custom-order sampling.
- Variable-size batches via a batch sampler.
- A `DistributedSampler` analog for data-parallel training (each rank sees a disjoint shard). `oversample`/`undersample`/`splitobs`/`shuffleobs` exist but aren't wired into `DataLoader`.
### B. `worker_init_fn` + per-worker RNG seeding
PyTorch runs a user function per worker at startup and auto-seeds each worker's RNG (`base_seed + worker_id`) so augmentation inside `getobs` differs per worker yet stays reproducible; `get_worker_info()` exposes the worker id. MLUtils has no per-worker init hook or per-worker seeding (relevant to the `num_workers` process path).
### C. Order-preserving iteration (PyTorch default; `in_order`)
PyTorch returns batches in order by default (v2.6 added `in_order=False` to opt out for throughput). MLUtils' `parallel` and `num_workers` **always** break ordering with no opt-in to preserve it. Missing an ordered mode (reorder worker results by batch index) for reproducible/eval pipelines.
### D. `pin_memory` (+ `pin_memory_device`)
PyTorch copies batches into page-locked host memory for faster async host→GPU transfer. MLUtils has nothing, by design (framework/GPU-agnostic, no CUDA dep). Likely **out of scope for MLUtils core** — belongs in a GPU-aware wrapper (Flux / MLDataDevices). Listed for completeness.
### E. `timeout` (worker watchdog)
PyTorch errors if a batch isn't produced within N seconds, catching hung/deadlocked workers. MLUtils' `remotecall_fetch` waits indefinitely, so a stuck worker hangs the whole loader.
### F. Iterable-style datasets (streaming / unknown length)
PyTorch supports `IterableDataset` (`__iter__`, no `__len__`) for streaming/infinite sources, sharded across workers. MLUtils is **map-style only** — requires finite, indexable `numobs`/`getobs`. No streaming support; would need a separate iteration path.
### G. `persistent_workers` toggle (inverse gap)
PyTorch tears workers down each epoch by default; `persistent_workers=True` keeps them. MLUtils is effectively *always* persistent (global warm pool) — the inverse gap (no per-epoch teardown).
---
Related distributed-loading follow-ups (from the `num_workers` design, #240): bring-your-own `AbstractWorkerPool`, a `prefetch`/`channelsize` kwarg (PyTorch `prefetch_factor`), ~~a bug where concurrent/nested loaders share and contend for the same pooled workers~~ (now fixed: each live loader leases a disjoint block of workers), and ~zero-copy/shared-memory IPC~ (done in #244).
Contributor guide
No contributing guide indexed for this repository
Research direction
Start with the MLUtils.DataLoader interface and compare its listed options with the PyTorch DataLoader features in sections A–G. Read the related num_workers design in issue #240 and separate in-scope work from items marked out of scope or already completed. Done requires choosing one narrowly defined feature and documenting its tests and acceptance criteria.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- julia, pytorch
- Domain
- data, machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Needs clarification
- Newbie friendliness
- 25/100