DataLoader: batching, shuffling and workers

What DataLoader adds

DataLoader(dataset, batch_size=32, shuffle=True) turns single samples into batches. Its default collate function stacks matching pieces: 32 samples of (x of shape (3,), y) become x of shape (32, 3) and y of shape (32,). With shuffle=True the order changes every epoch, which stops the model learning from the order of the data. Use shuffle=False for validation and test, where order does not matter and reproducibility helps.

num_workers=4 starts four worker processes that call __getitem__ in parallel while the main process trains, which matters when loading or augmenting images is slow. On Windows and macOS (and on Linux from Python 3.14, which defaults to forkserver), workers are started in fresh interpreters that re-import your script, so the entry code must sit under if __name__ == "__main__": or it re-runs in each worker. pin_memory=True speeds up copies to a CUDA GPU. drop_last=True discards a final short batch, useful when a layer such as BatchNorm misbehaves on tiny batches.

When samples have different lengths (sentences, audio clips), default stacking fails; supply a collate_fn that pads them.

Batches from a TensorDataset

100 samples with batch_size=32 give three full batches of 32 and a final batch of 4. Each xb has shape (32, 3) and each yb has shape (32,), except the last pair with 4.

import torch
from torch.utils.data import DataLoader, TensorDataset

X = torch.randn(100, 3)
y = torch.randint(0, 2, (100,))
ds = TensorDataset(X, y)

loader = DataLoader(ds, batch_size=32, shuffle=True)
for xb, yb in loader:
    print(xb.shape, yb.shape)

Padding variable-length sequences

pad_sequence pads shorter sequences with zeros. Returning the true lengths lets the model ignore the padding later.

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

def pad_collate(batch):
    seqs, labels = zip(*batch)               # tuples of tensors
    lengths = torch.tensor([len(s) for s in seqs])
    padded = pad_sequence(seqs, batch_first=True, padding_value=0)
    return padded, torch.stack(labels), lengths

batch = [(torch.tensor([5, 2, 9]), torch.tensor(1)),
         (torch.tensor([7, 3]), torch.tensor(0))]
x, y, lengths = pad_collate(batch)
print(x)          # [[5, 2, 9], [7, 3, 0]]
print(lengths)    # [3, 2]
# Use it with: DataLoader(ds, batch_size=32, collate_fn=pad_collate)

Dabbawalas and lunchboxes

Mumbai's dabbawalas do not carry one tiffin at a time. Several collectors (workers) gather boxes in parallel, group them by destination into crates (batches), and hand each crate to the train (the GPU) in one go. If the collectors are too few, the train waits at the platform; that idle train is a GPU starved by a slow data pipeline.

Quick check: What does `num_workers=4` do in a `DataLoader`?

• Trains the model on four GPUs
• Loads samples in four background processes while the main process trains
• Splits each batch into four smaller batches
• Repeats each epoch four times