JuliaML / JuliaML/MLUtils.jl

Tracking: DataLoader feature parity with PyTorch

Open
#241 0 comments 0 reactions 0 assignees View on GitHub
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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.