DataLoader and Batching

What DataLoader does

DataLoader wraps a Dataset and adds: batching (group N examples into a tensor of shape [N, ...]), shuffling (sampling order), parallel loading (multiple worker processes), and pin memory (for faster CPU→GPU transfer). One line at the top of your training loop, and the rest of the code can just iterate over batches.

The two arguments that matter most are batch_size (tune for VRAM and throughput) and num_workers (tune for CPU/IO bandwidth). Default num_workers=0 is single-process, which is fine for tiny datasets and a debugging nightmare for everything else.

Tip: Set num_workers to roughly the number of physical CPU cores you have available. On a 16-core box, num_workers=8 is a sensible starting point. Test with htop open — workers should be busy, GPU should be busy, neither should be idle waiting on the other.

Common pitfalls

Forgetting shuffle=True on training — if your data is class-sorted, the model overfits the first class and never recovers. Always shuffle the training loader.

Shuffling the validation loader — wastes time and breaks reproducibility of per-batch metrics. Validation loaders use shuffle=False.

Using num_workers>0 on CUDA without pin_memory=True — leaves CPU→GPU bandwidth on the table. Always pair them.

Mutating shared state in __getitem__ — workers are subprocesses, so any global state is per-worker. Logging counters, RNG state, etc. behave subtly differently than you expect.

Custom collate

The default collate_fn stacks tensors with the same shape. For variable-length sequences (text, audio, sets), you need a custom collate that pads or packs. Hugging Face's DataCollatorWithPadding handles the common case for you.

Principle: Three things to set on every training DataLoader: shuffle=True, num_workers>0, pin_memory=True on CUDA. Three things to set on every validation DataLoader: shuffle=False, num_workers>0, pin_memory=True.

Code

Train and val DataLoaders, the right way·python

from torch.utils.data import DataLoader

train_loader = DataLoader(
    train_ds,
    batch_size=128,
    shuffle=True,
    num_workers=8,
    pin_memory=True,
    persistent_workers=True,   # avoid worker spawn cost each epoch
    drop_last=True,            # drop the last partial batch (cleaner stats)
)

val_loader = DataLoader(
    val_ds,
    batch_size=256,            # bigger batch for inference
    shuffle=False,
    num_workers=4,
    pin_memory=True,
    persistent_workers=True,
)

Custom collate for variable-length sequences·python

import torch
from torch.nn.utils.rnn import pad_sequence

def pad_collate(batch):
    seqs, labels = zip(*batch)        # tuples of variable-length tensors and scalars
    seqs = pad_sequence(seqs, batch_first=True, padding_value=0)
    labels = torch.tensor(labels)
    return seqs, labels

loader = DataLoader(my_dataset, batch_size=32, collate_fn=pad_collate)

What DataLoader does

Common pitfalls

Custom collate

Code

External links

Exercise

Progress

Comments 0