Datasets, DataLoaders and Collate: Coverage Before Throughput

The story — find the waiting stage

A slow training step can be limited by storage, decoding, collation, transfer, launch overhead or the model. The original Tesla/V100 anecdote and weeks-to-days figures were not verified measurements. The useful lesson is to isolate stages before tuning workers: low GPU utilization alone does not identify the data loader as the culprit.

A typical map-style pipeline separates pick (sampler), fetch (Dataset) and stack (collate). DataLoader coordinates these jobs, batching and optional workers. Iterable sources use a different sampling contract. Neither design requires downloading anything: all executed examples here use local synthetic records or tiny generated tar fixtures.

Prerequisites and outcomes

Bring tensor semantics and Module ownership. Runtime: PyTorch2.9.0+cpu. You will write indexed and iterable datasets, test sample coverage, collate ragged sequences with lengths/masks, reason about worker replicas and prefetch, and specify checkpoint/resume state. Optional pre-reading: custom datasets, loading talk, official tutorial, PyTorch2.9 data reference.

1 · Pick, fetch and stack

A map-style Dataset normally implements __getitem__ and __len__; a sampler supplies keys/indices. Batched __getitems__ can optionally optimize bulk retrieval. DataLoader applies collation when automatic batching is enabled. A Dataset can return prebatched data when that is the right design; the point is separation of responsibility, not prohibition.

import torch
import numpy as np
import time
import random
import io
import tarfile
from pathlib import Path
from torch.utils.data import Dataset,DataLoader,IterableDataset,get_worker_info
from torch.nn.utils.rnn import pad_sequence
 
torch.set_num_threads(1)
 
class Records(Dataset):
    def __init__(self,n=23):self.n=n
    def __len__(self):return self.n
    def __getitem__(self,i):
        if not 0<=i<self.n:raise IndexError(i)
        return torch.tensor([float(i),float(i*i)]),i
 
ds=Records()
loader=DataLoader(ds,batch_size=7,shuffle=True,num_workers=0,
                  generator=torch.Generator().manual_seed(16))
seen=[]
for xb,yb in loader:
    assert torch.equal(xb[:,0].long(),yb)
    seen.extend(yb.tolist())
assert sorted(seen)==list(range(23)) and len(seen)==len(set(seen))
assert [len(y) for _,y in DataLoader(ds,batch_size=7)]==[7,7,7,2]
assert sum(len(y) for _,y in DataLoader(ds,batch_size=7,drop_last=True))==21

The warehouse/truck analogy maps retrieval versus scheduling. It does not imply all pipelines have exactly two objects or that a worker count makes arbitrary state safe. This ID-based test catches a dataset that ignores i and returns one debug record repeatedly, something shuffle=True cannot repair.

With automatic batching and positive workers, collation normally runs in workers; one batch is not necessarily assembled by combining samples from every worker. With zero workers these stages run in the main process. Pinning is a host-memory stage, not the device transfer itself. The diagram shows the conceptual pipeline; only CPU loading is executed here.

2 · MNIST, CIFAR and existing local files

For an already installed compatible torchvision and already local data, the documented recipes are:

from torchvision import datasets, transforms
train_ds = datasets.MNIST(root, train=True, download=False,
                          transform=transforms.ToTensor())
train_loader = DataLoader(train_ds, batch_size=128, shuffle=True, num_workers=0)
cifar_ds = datasets.CIFAR10(root, train=True, download=False,
                            transform=transforms.ToTensor())

This is an unexecuted local-data adapter, shown as text so copying all executable fences never downloads/imports an absent optional library. Missing data raises. MNIST items are28×28 grayscale; CIFAR10 is32×32 RGB. ToTensor converts supported uint8 images to channel-first float values scaled to[0,1]; arbitrary float arrays are not always divided by255. Dataset-specific transforms, splits and channels still matter. ImageNet's data acquisition/structure is not identical to an automatically downloadable toy dataset.

DL011 supplies a complete NumPy local IDX reader and training recipe. Its normalized arrays can be wrapped with TensorDataset and fed to DL015's training/evaluation functions. This maintains a complete local-only MNIST route without claiming we trained on those digits.

3 · A custom dataset with explicit file ownership

Enumerate class directories and filenames deterministically, filter regular files, record the class-to-index mapping and reuse it across splits. Decode inside a context manager so a file handle is not retained across worker creation. The optional image adapter below imports Pillow only when an image is actually requested; its source review is not an assertion that Pillow/image decoding was executed in this runtime.

class ImageFolderDataset(Dataset):
    def __init__(self,root,class_to_idx=None,transform=None):
        root=Path(root)
        names=sorted(p.name for p in root.iterdir() if p.is_dir())
        self.class_to_idx=class_to_idx if class_to_idx is not None else {c:i for i,c in enumerate(names)}
        self.transform=transform
        self.samples=[]
        for c in names:
            if c not in self.class_to_idx:raise ValueError('unknown class '+c)
            for p in sorted((root/c).iterdir()):
                if p.is_file() and p.suffix.lower() in {'.png','.jpg','.jpeg','.ppm'}:
                    self.samples.append((p,self.class_to_idx[c]))
        if not self.samples:raise ValueError('no supported images')
    def __len__(self):return len(self.samples)
    def __getitem__(self,i):
        from PIL import Image
        path,label=self.samples[i]
        with Image.open(path) as src:
            image=src.convert('RGB')
            if self.transform is not None:image=self.transform(image)
            else:image=torch.from_numpy(np.asarray(image,dtype=np.uint8).copy()).permute(2,0,1).float()/255
        return image,label

Images must have equal transformed shape for default tensor stacking; add an explicit resize/crop policy when needed. The toy Records class in §1 is a fully executed custom Dataset. For the original “20 local files” exercise, create a manifest with path,label,split, validate all labels, then fetch several different IDs and check one batch. Do not read arbitrary personal files merely to manufacture an example.

Dataset initialization should avoid unnecessary expensive decoding, but metadata scanning can itself be substantial. __getitem__ need not be mathematically pure: seeded random augmentation is common, and a worker-local cache/connection can be valid. Avoid unsafe shared mutable state and untracked side effects. Numeric NumPy arrays are normally supported by default_collate; raw PIL images are not. Fork can inherit file descriptors, but ownership/library safety is not guaranteed; spawn requires serializable state. Open database/file resources within each worker and close them appropriately rather than claiming handles never survive fork.

4 · DataLoader options and their constraints

4.1 batch_size

This determines automatic batch grouping; final size may be smaller. Increasing it changes activation/prefetch memory and optimization updates, not all memory linearly because parameters and other costs stay fixed. With23 records and batch7, four batches include a final two-row batch. Weight epoch metrics by examples, not batch count.

4.2 shuffle

For map-style datasets, shuffle=True selects randomized sampling; the generator's state advances across iterations. Use one seeded generator without resetting it each epoch when you want reproducible but changing order. Ordered validation is convenient, but an order-invariant properly weighted metric can remain valid under shuffling. Time-series/order-sensitive tasks and custom samplers can require different policies. In distributed training call DistributedSampler.set_epoch(epoch) to vary epoch order, and configure ranks correctly; this does not by itself repair overlapping datasets.

4.3 num_workers

Zero fetches in the main process and gives clearer tracebacks. Positive counts use subprocesses. Start at zero, compare modest counts and measure loader wait, CPU, resident memory and storage throughput. More is not necessarily faster; copies/serialization, CPU oversubscription and I/O contention can dominate. Process-start behavior depends on OS/Python configuration—do not assume fork everywhere. Use top-level picklable definitions and a main guard for spawn.

4.4 pin_memory

Pinned host buffers can improve host-to-device transfer. non_blocking=True requests asynchronous behavior where possible; pinned memory and non-blocking flags do not alone guarantee copy/compute overlap. Useful overlap also needs suitable hardware, stream scheduling and independent work. Non-blocking copies can have benefits even without pinning in some paths. Pinning consumes limited host resources; it is not “always free” or a meaningful accelerator optimization for this CPU-only test.

4.5 drop_last

For map-style automatic batching it discards the final incomplete batch. This can help fixed-size assumptions but reduces sample coverage; keep validation coverage unless explicitly accounting for exclusions. BatchNorm's failure depends on the total elements per channel, not only batch dimension in all variants. For multi-worker IterableDataset, dropping can occur per replica, losing more than one partial batch.

4.6 persistent_workers

It requires a positive worker count. Keeping processes/dataset replicas alive avoids repeated worker startup, not necessarily re-running a Dataset constructor every epoch. Worker-local state and RNG streams also persist; a mutation of the main-process dataset may not propagate. Use it only when its lifetime semantics are intended.

4.7 prefetch_factor

For positive workers, the documented default is two batches per worker. A rough queued-batch payload estimate is workers*prefetch_factor*batch_bytes, plus worker memory, in-flight and pinned copies. With zero workers, do not set a multiprocessing prefetch factor. Extra prefetch can hide latency but cannot exceed sustained storage/decode throughput and can increase waste on restart.

5 · Collation for variable-length sequences

Default collate handles supported nested structures and numeric arrays/tensors. Raw strings can remain strings; “all text immediately fails” is false. Unequal tensor lengths cannot be directly stacked. Padding creates a rectangle but must be accompanied by lengths/masks and, for token losses, ignored targets or a selected-token denominator.

5.1 Reproduce the ragged-stack failure

sentences=[torch.tensor([1,2,3]),torch.tensor([4,5,6,7,8]),torch.tensor([9])]
try:torch.stack(sentences)
except RuntimeError:pass
else:raise AssertionError('ragged tensors cannot stack directly')

5.2 Padding, lengths and masks

A mask derived from token!=0 is wrong when zero is a valid token. Derive it from lengths. Define the empty-sequence policy: this collator permits individual empties but rejects an all-empty batch, leaving at least one valid token overall. Sequence-classification labels have shape B; a language-model target/loss needs an additional token-level contract.

def collate_text(batch):
    if not batch:raise ValueError('empty batch')
    tokens=[torch.as_tensor(x,dtype=torch.long) for x,_ in batch]
    if any(x.ndim!=1 for x in tokens):raise ValueError('tokens must be vectors')
    lengths=torch.tensor([len(x) for x in tokens],dtype=torch.long)
    if lengths.max().item()==0:raise ValueError('all sequences empty')
    labels=torch.tensor([y for _,y in batch],dtype=torch.long)
    padded=pad_sequence(tokens,batch_first=True,padding_value=0)
    mask=torch.arange(padded.shape[1])[None,:]<lengths[:,None]
    return {'tokens':padded,'lengths':lengths,'mask':mask,'labels':labels}
 
batch=collate_text([(torch.tensor([0,2,3]),1),(torch.tensor([4]),0)])
assert batch['tokens'].tolist()==[[0,2,3],[4,0,0]]
assert batch['mask'].tolist()==[[True,True,True],[True,False,False]]
assert batch['lengths'].tolist()==[3,1]
assert batch['mask'].sum().item()==4
# Masked token-loss accounting, not a sequence-classification loss.
losses=torch.tensor([[1.,2.,3.],[4.,100.,100.]])
masked=(losses*batch['mask']).sum()/batch['mask'].sum()
assert masked.item()==2.5 and losses.mean().item()!=2.5
try:collate_text([(torch.tensor([],dtype=torch.long),0)])
except ValueError:pass
else:raise AssertionError('all-empty token denominator must reject')

5.3 Bucketing by length

Group similar lengths to reduce padding, then shuffle buckets/batch order according to a documented policy. For lengths1,2,9,10 at batch2, grouping (1,2),(9,10) pads to4+20=24 positions rather than20+18=38 for (1,10),(2,9), while valid tokens total22 in either case. A custom BatchSampler emits index lists; do not simultaneously pass incompatible batch_size/shuffle/sampler settings. Bucketing changes order and resumability state, not labels or the loss's valid-token denominator.

6 · Diagnose slow training in four stages

1. Model-only baseline: reuse a resident batch, warm up, time forward/backward/update. On CUDA synchronize or use events appropriately; otherwise CPU timing can measure launch latency rather than work.
2. Loader-only/end-to-end: time iteration with workers0 then modest positive counts; record first-batch versus steady-state wait. A speedup suggests pipeline work matters but does not uniquely identify disk versus decode versus IPC.
3. Profile fetch/transform/collate: inspect individual stages and memory. WebP is not universally faster than JPEG, and moving augmentation to a GPU changes resource use rather than providing a free improvement.
4. Transfer/scheduling: inspect dtype/layout, pinning and nonblocking copies, then profile actual overlap. Do not claim low utilization plus a flag proves the bottleneck solved.

The CPU timing fixture is deterministic arithmetic, not an image-decoding or GPU benchmark. Test0/1/2 workers at most and run it only from a script main guard; the review records the actual results without assuming the curve must flatten or reverse.

def benchmark_workers():
    results={}
    for workers in (0,1,2):
        loader=DataLoader(Records(512),batch_size=32,num_workers=workers)
        start=time.perf_counter();ids=[]
        for _,labels in loader:ids.extend(labels.tolist())
        assert ids==list(range(512))
        results[workers]=time.perf_counter()-start
    return results

7 · Failure cases: ownership, identity and memory

A SQLite connection created in one process should not be reused as if it were a safe process-independent handle. Open lazily per worker or in worker initialization, using the library's process-safety guidance. The original SQLite thread-error string did not prove a particular fork failure. File-format decoders, HDF5 and database clients have different constraints; “same for every handle” is too broad.

A Dataset that always returns index0 defeats any sampler. The ID coverage assertions in §1 catch missing/duplicate rows. Worker replicas may multiply Python-object memory; NumPy backing arrays, copy-on-write, shared tensors and memory maps have different residency behavior. Do not calculate every workload as exactly workers×parentRAM. Inspect resident/proportional memory and avoid enormous Python lists of objects when an indexed array/manifest suffices.

8.5 · IterableDataset, local tar streams and sharding

Dataset size alone does not prohibit map-style access: disk-backed indices and memory maps can support enormous corpora. Choose IterableDataset when sequential access is natural or random access unavailable. Every worker receives a replica; partition across distributed rank and worker, or duplicates result. A bounded shuffle buffer is not a uniform permutation of an arbitrarily long stream.

This complete local tar reader expects regular .npy members named <label>-<id>.npy, with numeric3D (C,H,W) arrays. It reads members without extracting paths to disk, rejects oversize members and unsafe pickled NumPy objects, and limits shuffle-buffer count. It does not contact S3.

def assigned_shards(paths,rank=0,world=1,worker=0,workers=1):
    if not (world>0 and workers>0 and 0<=rank<world and 0<=worker<workers):
        raise ValueError('invalid rank/worker partition')
    return list(paths)[rank*workers+worker::world*workers]
 
def buffer_shuffle(items,size,seed):
    if size<1:raise ValueError('positive buffer size required')
    rng=random.Random(seed);buf=[]
    for item in items:
        if len(buf)<size:buf.append(item);continue
        j=rng.randrange(len(buf));yield buf[j];buf[j]=item
    rng.shuffle(buf);yield from buf
 
class ShardedTarDataset(IterableDataset):
    def __init__(self,paths,rank=0,world=1,shuffle_size=1,seed=0):
        super().__init__();self.paths=list(paths);self.rank=rank;self.world=world
        self.shuffle_size=shuffle_size;self.seed=seed
    def __iter__(self):
        info=get_worker_info();wid=info.id if info else 0;nw=info.num_workers if info else 1
        paths=assigned_shards(self.paths,self.rank,self.world,wid,nw)
        def records():
            for path in paths:
                with tarfile.open(path,mode='r|*') as tar:
                    for member in tar:
                        if not member.isfile() or not member.name.endswith('.npy'):continue
                        if member.size>4_000_000:raise ValueError('oversize record')
                        label=int(Path(member.name).name.split('-',1)[0])
                        stream=tar.extractfile(member)
                        if stream is None:raise ValueError('missing regular member')
                        with stream:raw=stream.read()
                        arr=np.load(io.BytesIO(raw),allow_pickle=False)
                        if arr.ndim!=3 or arr.dtype.kind not in 'uif':raise ValueError('image array contract')
                        if not np.isfinite(arr).all():raise ValueError('nonfinite image')
                        yield torch.from_numpy(arr.astype(np.float32,copy=True)),label
        yield from buffer_shuffle(records(),self.shuffle_size,self.seed+self.rank*nw+wid)
 
paths=list(range(17))
parts=[assigned_shards(paths,r,2,w,3) for r in range(2) for w in range(3)]
flat=[x for part in parts for x in part]
assert sorted(flat)==paths and len(flat)==len(set(flat))
assert sorted(buffer_shuffle(range(20),4,3))==list(range(20))
assert list(buffer_shuffle(range(20),4,3))==list(buffer_shuffle(range(20),4,3))

The partition assumes equal worker counts across ranks and an identical ordered shard manifest. Shard counts/workloads can be uneven; few large shards may starve workers. A real resume protocol must also handle buffered samples and per-worker positions. A10,000-item shuffle buffer can be large for decoded images; size it in bytes as well as counts.

S3 stretch recipe, unexecuted: choose an existing authenticated streaming adapter, obtain a binary file-like response per selected shard, and pass it as tarfile.open(fileobj=response,mode='r|gz') within the response's context manager. Keep the same rank/worker partition and bounded decoder; define retries so a partially delivered shard is not silently replayed. Record object version/ETag, shard order, member cursor, shuffle RNG/buffer and consumer position. No cloud credentials, S3 requests or new infrastructure are needed for the local counterpart above. A WebDataset reader can provide transport/format utilities, but its sharding and duplicate/resume behavior still need explicit tests.

8.6 · StatefulDataLoader and exact resume boundaries

The tagged TorchData v0.8.0 source contains the reviewed StatefulDataLoader path; it is a separate package, not an automatic feature of every torch.utils.data.DataLoader. It offers state_dict/load_state_dict; custom stateful datasets/samplers can participate, and iterable worker count and cursor semantics matter. TorchData's later deprecation/removal of its DataPipes/DataLoader2 direction does not mean every torch.utils.data.IterDataPipe API vanished from core PyTorch.

from torchdata.stateful_dataloader import StatefulDataLoader
loader = StatefulDataLoader(dataset, batch_size=128, num_workers=4)
checkpoint['loader'] = loader.state_dict()
# Reconstruct compatible dataset/loader and restore, together with model/optimizer/RNG:
loader.load_state_dict(checkpoint['loader'])

This optional-package adapter is source-reviewed, not imported/executed here. Exact replay must be demonstrated: save a mid-epoch state, record the next several IDs/batches, reconstruct in a fresh process, restore and compare. A random seed alone does not identify a prefetched stream's consumed position. A simple local cursor model below illustrates the contract without pretending to implement multiprocessing loader state.

class CursorStream:
    def __init__(self,items):self.items=list(items);self.pos=0
    def __iter__(self):
        while self.pos<len(self.items):
            item=self.items[self.pos];self.pos+=1;yield item
    def state_dict(self):return {'pos':self.pos,'items':self.items.copy()}
    def load_state_dict(self,state):
        if state['items']!=self.items or not 0<=state['pos']<=len(self.items):raise ValueError('cursor mismatch')
        self.pos=state['pos']
a=CursorStream(range(9));it=iter(a);assert [next(it) for _ in range(4)]==[0,1,2,3]
saved=a.state_dict();tail=list(it)
b=CursorStream(range(9));b.load_state_dict(saved)
assert list(b)==tail==[4,5,6,7,8]

8.7 · Hugging Face datasets and storage choices

The versioned datasets documentation distinguishes an Arrow-backed random-access Dataset from a streaming IterableDataset. load_dataset(...,streaming=True) can start iterating without a complete download, but iteration still reads remote data; streaming is not offline. with_format('torch') formats supported columns, not arbitrary un-tokenized strings into a model-ready padded tensor.

For local text/tabular records use already-available in-memory/Arrow constructors or a local manifest; choose memory mapping for random access beyond RAM and sequential shards for streaming. Dataset size thresholds such as10GB/100GB or “90% of all projects” do not determine the right API. Measure storage layout, random-access needs, transform cost, worker duplication, shuffle quality and resume requirements. WebDataset/Hugging Face tooling can help but does not remove these contracts.

Fully resident data minimizes I/O when it fits; memory mapping lets the OS page files and share backing storage but random access can still thrash; streaming bounds buffering by design but weakens global shuffle and may lack a meaningful length/epoch. Buffered decoded tensors and worker queues mean streaming is not literally zero or constant memory under every configuration.

9 · Retention scaffold — six worked answers

1. Interfaces: map-style normally defines len/getitem; IterableDataset defines iter and may lack length. Optional bulk access/prebatched modes exist.
2. Initialization/state: keep metadata manageable and decoding lazy where useful; random transforms and worker-local caches are valid with deliberate RNG/ownership. Avoid sharing unsafe connections across process starts.
3. Pin/nonblocking: can improve transfer and host scheduling, but either may help separately and together do not prove copy/compute overlap. Streams, dependencies and hardware decide.
4. Collate: ragged tensors cannot stack; pad to batch maximum, return lengths/mask and define target-loss masking. A valid zero token must not be mistaken for padding.
5. Diagnosis: resident model baseline; loader-only/end-to-end timing; fetch/transform/collate profiling; transfer/overlap profiling. Synchronize GPU measurements and distinguish first-batch startup.
6. Streaming stretch: §8.5 provides a complete local tar decoder, rank+worker partition and bounded shuffle. The S3 adapter/resume procedure is explicit and unexecuted; the review tests local shards, coverage and malformed data, not cloud availability.

Reconstruct the masked mean2.5 and uneven final batch tomorrow. Next week verify cursor continuity and ask which queued state a real restart must save. The sticky rule is pick, fetch, stack, measure—not “low utilization means blame data.”

Sources and execution boundary

PyTorch2.9 data, pin/nonblocking tutorial, torchvision v0.24 MNIST, CIFAR, TorchData v0.8 stateful loader, TorchData v0.9 README, datasets v4.2 streaming, WebDataset. Exact passages/hashes distinguish versioned facts from mutable guidance. Local synthetic CPU loading/collation/sharding contracts were tested; no external dataset, S3, CUDA transfer, optional-package install or real-model training result is claimed.