PyTorch DataLoader: batches, shuffling, and worker processes

A PyTorch DataLoader supplies batches of examples to your training or evaluation loop. The dataset defines the examples; the loader controls how they are collected and delivered. Before tuning its speed, check that each batch has the shape, type, and sample coverage your model expects.

This guide starts with a small CPU example you can run without downloading a dataset. It covers batching, shuffling, worker processes, and variable-length inputs, with checks that make missing or repeated samples visible.

Dataset and DataLoader: different responsibilities

A map-style Dataset defines __len__ and __getitem__: how many examples exist and how to retrieve one by index. A DataLoader groups the retrieved examples into batches. The official datasets tutorial introduces this separation.

In the example below, each example is a dictionary containing three features, one class label, and a sample ID. The ID is a diagnostic field. Keeping it available during a small trial lets us verify coverage independently of the model's loss.

Use an environment with PyTorch installed and record its version before comparing results. Our development environment guide explains how to keep project dependencies separate.

A runnable PyTorch DataLoader example

Save this as dataloader_demo.py. It creates nine synthetic examples and requests batches of four. No GPU, model, or external data files are required.

import argparse
import torch
from torch.utils.data import DataLoader, Dataset


class ToyDataset(Dataset):
    def __len__(self):
        return 9

    def __getitem__(self, index):
        return {
            'features': torch.tensor([index, index + 1, index + 2],
                                     dtype=torch.float32),
            'label': torch.tensor(index % 2, dtype=torch.long),
            'sample_id': index,
        }


def collect(workers=0, shuffle=True, drop_last=False, seed=11):
    loader = DataLoader(
        ToyDataset(), batch_size=4, shuffle=shuffle,
        num_workers=workers, drop_last=drop_last,
        generator=torch.Generator().manual_seed(seed),
    )
    sizes, ids = [], []
    for batch in loader:
        x, y = batch['features'], batch['label']
        assert x.shape == (len(y), 3)
        assert x.dtype == torch.float32 and y.dtype == torch.long
        sizes.append(len(y))
        ids.extend(batch['sample_id'].tolist())
    return sizes, ids


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--workers', type=int, default=0)
    args = parser.parse_args()
    sizes, ids = collect(workers=args.workers)
    assert sorted(ids) == list(range(9))
    print('batch sizes:', sizes)
    print('samples:', len(ids), 'unique:', len(set(ids)))

Run the file first with loading in the main process, then with two worker processes:

python dataloader_demo.py --workers 0
python dataloader_demo.py --workers 2

Both commands produced this output in our test:

batch sizes: [4, 4, 1]
samples: 9 unique: 9

The two full batches contain four examples each; the final batch contains one. Across the three batches, all nine sample IDs appear exactly once. The shape assertion checks that each feature tensor has three columns and the same number of rows as there are labels. It also checks the floating-point feature type and integer label type.

Those choices suit this toy classification dataset. Your own task may need images, token IDs, regression targets, or additional masks. Inspect a real batch before attaching the model instead of copying these dimensions into unrelated training code.

Understand batch size and the last batch

We tested the same loader with drop_last=True. It produced two batches of four and eight unique sample IDs. With shuffling disabled, ID 8 was omitted; with the seed and shuffling used here, ID 3 was omitted. That difference is a useful reminder to inspect the actual samples, not just the number of batches.

Keep the final partial batch when all examples should contribute. If your training design requires equal batch sizes, make the omission explicit and account for it in your run notes. For evaluation, verify full coverage and calculate aggregate metrics from their correct denominators. An unweighted average of a four-example batch mean and a one-example batch mean would give the last example too much influence.

Changing batch size also changes the number of iterations. In this nine-example test, batches of four yield three iterations without dropping the last batch. If a scheduler, logger, or checkpoint interval counts iterations, review that behavior when you change the batch size.

Shuffle training data without mixing your splits

The example enables shuffle=True and supplies a seeded generator. Our tests confirmed that the shuffled run retained all nine IDs when drop_last=False. Shuffling changed their order, not their membership.

Define training, validation, and test membership separately before building the loaders. A shuffled loader does not establish a valid split. For time series, grouped observations, or multiple records from one person, choose a split that respects those relationships. Keep a small record of sample IDs per split so accidental overlap is detectable.

A fresh loader with a fresh generator seeded to 11 repeated our initial order. Reusing a generator advances its state. Random transforms can need additional seeding, including Python or NumPy randomness through worker_init_fn. Follow PyTorch's reproducibility guidance; the same seed is not a guarantee of identical results across releases or platforms.

Choose num_workers after the loader works

num_workers=0 loads data in the main process. A positive value uses subprocesses. Start at zero while debugging, then compare a small number of workers on representative inputs. The PyTorch tuning guide describes worker-based loading and pinned memory.

Our synthetic example performs almost no input processing. Its successful two-worker run proves that the example works with that process arrangement; it does not establish a speed advantage. A real trial should include your file reads, decoding, transformations, and training loop.

Record elapsed time, examples processed, CPU use, and host-memory use for the same workload. Compare more than one run and separate startup from steady-state behavior. If extra workers increase memory use without reducing total time, keep the simpler configuration. A fast GPU cannot resolve a slow data source by itself.

Run worker examples safely on Windows

Windows uses spawned processes. Keep the dataset class and any custom collation function at module level, and start execution under if __name__ == '__main__':. The PyTorch Windows FAQ explains the main guard and recommends returning to zero workers when diagnosing worker failures.

The example follows that structure. If it works as a saved file but fails in a notebook, first reproduce the issue in the standalone file with the same interpreter. Keep the full traceback and reduce the dataset to the smallest failing example.

Use collate_fn for variable-length inputs

Our first dataset returns fixed-size feature vectors. For sequences with different lengths, define how to build a rectangular batch. This second example pads two sequences and returns their original lengths alongside their labels. Save it as padding_demo.py and run it with Python.

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


def pad_batch(samples):
    sequences, labels = zip(*samples)
    lengths = torch.tensor([len(x) for x in sequences])
    padded = pad_sequence(sequences, batch_first=True, padding_value=0)
    return padded, torch.tensor(labels, dtype=torch.long), lengths


if __name__ == '__main__':
    samples = [(torch.tensor([4, 8, 2]), 1), (torch.tensor([5]), 0)]
    padded, labels, lengths = pad_batch(samples)
    print(padded.tolist(), labels.tolist(), lengths.tolist())

The padded values are [[4, 8, 2], [5, 0, 0]], the labels are [1, 0], and the lengths are [3, 1]. PyTorch's pad_sequence reference documents the batch-first layout and padding value.

Pass collate_fn=pad_batch when creating a loader for a dataset that returns these sequence-label pairs. Our verification also used that function in a real two-worker DataLoader. Keep it at module level so worker processes can import it.

Padding creates positions that are not original observations. Use the returned lengths to construct the mask or sequence representation your model requires. Make sure the loss excludes padding where appropriate, and choose a padding value that fits your vocabulary. A numeric zero is only a placeholder in this example, not a universal padding-token rule.

Separate CPU loading from GPU transfer

pin_memory=True prepares suitable host-memory buffers; it does not move the batch onto a GPU. Device transfer remains a separate step. For CUDA workloads, evaluate it alongside transfers using non_blocking=True, following the PyTorch transfer tutorial. Benefits depend on the workload, and asynchronous transfers require care with when buffers are reused.

Keep the CPU example working before adding device placement. Our PyTorch on A100 guide covers CUDA setup and device checks. We did not test GPU transfer or measure GPU utilization for this article.

Tune advanced options one at a time

prefetch_factor counts batches prepared ahead per worker. persistent_workers=True keeps workers and their dataset instances alive between passes. Both can affect host-memory use and state lifetime. Check the DataLoader API for your version before changing them.

This walkthrough uses a map-style dataset in one training process. Streaming IterableDataset instances and distributed sampling need their own partitioning and coverage checks. Do not assume that adding workers or training ranks preserves this example's one-copy-per-ID behavior.

Checks before a long run

• Inspect batch shapes, types, labels, and any padding mask.
• Count sample IDs across a full pass and explain omissions or repetitions.
• Check split membership separately from sample order.
• Reproduce failures with zero workers and a small input set.
• Change one performance setting at a time and measure the complete workload.

Our ten CPU checks passed on Windows with Python 3.12.14 and PyTorch 2.15.0+cpu. They covered zero and two workers, shuffling on and off, retaining or dropping the last batch, repeated seeded order, and custom padding with two workers. These are correctness checks on synthetic data. They do not measure storage I/O, CUDA behavior, or training speed.