The DataLoader wraps a Dataset and handles batching, shuffling, and parallel data loading automatically — the object that actually feeds mini-batches into the training loop.
Basic Usage
from torch.utils.data import DataLoader
loader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=4)
for x_batch, y_batch in loader:
print(x_batch.shape) # (32, ...) -- individual examples automatically stacked into a batch
breakKey Parameters
| Parameter | Effect |
|---|---|
batch_size | How many examples per batch — see Mini-Batch Gradient Descent |
shuffle | Randomizes example order each epoch — True for training, typically False for validation/test |
num_workers | Number of parallel subprocesses loading data — speeds up data loading when it's a bottleneck relative to GPU compute |
drop_last | Whether to discard a final, smaller-than-usual batch when the dataset size isn't evenly divisible by batch_size |
Custom Batching With collate_fn
def custom_collate(batch):
# batch is a list of (feature, label) tuples from __getitem__
features = [item[0] for item in batch]
labels = [item[1] for item in batch]
padded_features = pad_sequences(features) # e.g. for variable-length sequences
return padded_features, torch.stack(labels)
loader = DataLoader(dataset, batch_size=32, collate_fn=custom_collate)A custom collate_fn is essential whenever individual examples can't simply be stacked directly into a batch — the most common case being variable-length sequences (text, audio) that need padding to a uniform length within each batch before they can form a single tensor.
Common Mistakes
- Setting
shuffle=Truefor validation or test loaders — shuffling matters for training (to avoid the model seeing data in a fixed, potentially biased order across epochs), but is unnecessary and can make debugging/comparison across runs harder for evaluation. - Setting
num_workerstoo high relative to the machine's actual CPU cores — this can create excessive overhead rather than genuinely speeding up loading.
Interview Relevance
Q: "Why might you need a custom collate_fn when working with variable-length text sequences?" The default collation simply stacks individual examples into a batch tensor, which requires every example to already have the same shape. Variable-length sequences don't naturally have this property — a custom collate_fn is needed to pad shorter sequences (typically with a designated padding token) up to the batch's maximum length before stacking them into one uniform tensor.
Practice Question
Why is shuffle=True important for a training DataLoader but generally unnecessary for a validation DataLoader?
