Datasets, DataLoaders, Batching, and Shuffling

PyTorch models do not train on files, rows, or images directly. They train on tensors delivered in a steady sequence of mini-batches. The purpose of Dataset and DataLoader is to separate data access from training logic: the dataset knows how to retrieve one example, while the loader knows how to group examples, optionally shuffle them, move work into worker processes, and hand batches to the loop.

By the end of this lesson, you should be able to build a map-style dataset, choose a batch size, explain what shuffling changes, write a custom collate_fn, and diagnose the most common loader failures. In this PyTorch course, this is the bridge between tensor fundamentals and real training loops: a correct model can still learn poorly if the data pipeline feeds biased orderings, mismatched shapes, or mislabeled targets.

Purpose and Outcome

A Dataset gives PyTorch a Python object with a stable indexing interface. For a map-style dataset, __len__ reports how many examples are available and __getitem__(index) returns the example at one integer index. A DataLoader wraps that dataset and produces an iterator of batches. The training loop then becomes simple: iterate over batches, run the model, compute loss, backpropagate, and update parameters.

This separation matters because data loading has concerns the model should not own. A data pipeline may read JPEG files, tokenize text, normalize tabular columns, pad variable-length sequences, or sample minority classes more often. The model should receive tensors with predictable shapes and dtypes, not know where the rows came from.

How DataLoader Works Internally

At iteration time, DataLoader asks a sampler for indices. With shuffle=False, the default sampler yields indices in order. With shuffle=True, it uses a random sampler that visits each index once per epoch in a randomized order. The loader groups those indices according to batch_size, calls the dataset for each index, and passes the list of examples to a collation function.

The default collation function is more useful than it first appears. If each example is a tuple like (features, target), it transposes the list of examples into one batched feature tensor and one batched target tensor. If each example is a dictionary with identical keys, it collates each key independently. Numeric tensors are stacked along a new leading dimension, so four feature tensors shaped [3] become one tensor shaped [4, 3]. Strings and objects that cannot be stacked require more care.

When num_workers=0, loading happens in the main process. This is easiest to debug. When num_workers>0, worker processes fetch examples in parallel and send batches back to the main process. Parallel loading can hide disk or preprocessing latency, but it also requires that dataset state be picklable and that each worker avoid accidentally sharing unsafe resources such as open database cursors. pin_memory=True can speed transfer to CUDA because batches are allocated in page-locked host memory, but it is only useful when moving tensors to a GPU.

API Anatomy

The core constructor is DataLoader(dataset, batch_size=..., shuffle=..., sampler=..., collate_fn=..., drop_last=..., num_workers=...). Use batch_size to control how many examples are grouped together. Use shuffle=True for most training loaders, not validation or test loaders. Use drop_last=True when every batch must have exactly the same size, for example with some batch-normalization-sensitive experiments or distributed training setups. Use collate_fn when examples cannot be stacked by the default rule.

shuffle and sampler are alternatives. If you pass a custom sampler, it owns the index order, so PyTorch does not also accept shuffle=True. Weighted, grouped, distributed, and curriculum learning pipelines usually express their policy with a sampler.

Example 1: A Map-Style Dataset

This example builds a small tabular dataset. Each item returns a dictionary, and the default collator stacks matching keys into batched values. The output shows that the batch dimension is added in front of the three feature columns.

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

class ToyTabularDataset(Dataset):
    def __init__(self):
        self.x = torch.arange(12, dtype=torch.float32).reshape(4, 3)
        self.y = torch.tensor([0.0, 1.0, 0.0, 1.0])

    def __len__(self):
        return len(self.y)

    def __getitem__(self, index):
        return {"features": self.x[index], "target": self.y[index], "id": f"row-{index}"}

loader = DataLoader(ToyTabularDataset(), batch_size=2, shuffle=False)

for batch in loader:
    print(batch["features"].shape, batch["target"], batch["id"])

Expected behavior: two batches are produced. The feature shape is torch.Size([2, 3]) for each batch, targets are stacked into tensors, and row identifiers remain grouped with the matching examples.

Example 2: Variable-Length Items Need Collation

Text, audio, and event sequences often have different lengths. The default collator cannot stack [2], [1], and [3] tensors into one rectangular tensor. A custom collate_fn can pad sequences and return lengths so the model can ignore padding later.

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

examples = [
    (torch.tensor([4, 5]), 1),
    (torch.tensor([8]), 0),
    (torch.tensor([2, 2, 9]), 1),
]

def pad_collate(batch):
    tokens, labels = zip(*batch)
    lengths = torch.tensor([len(item) for item in tokens])
    padded = pad_sequence(tokens, batch_first=True, padding_value=0)
    return padded, lengths, torch.tensor(labels)

loader = DataLoader(examples, batch_size=2, collate_fn=pad_collate, shuffle=False)

tokens, lengths, labels = next(iter(loader))
print(tokens)
print(lengths)
print(labels)

The first batch contains two sequences. The shorter one is padded with zero, so the printed token tensor is rectangular. The lengths tensor tells later code that the real sequence lengths are two and one, while labels remain aligned with their original examples.

Example 3: Training Batches and Shuffle Policy

For stochastic gradient descent, batch order affects optimization. Shuffling reduces the chance that the model sees all easy, hard, class-zero, or class-one examples in long consecutive runs. The dataset itself remains unchanged; only the order of sampled indices changes for an epoch.

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

torch.manual_seed(0)
features = torch.randn(20, 4)
targets = (features.sum(dim=1, keepdim=True) > 0).float()
loader = DataLoader(TensorDataset(features, targets), batch_size=5, shuffle=True)
model = nn.Linear(4, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
loss_fn = nn.BCEWithLogitsLoss()

for batch_features, batch_targets in loader:
    logits = model(batch_features)
    loss = loss_fn(logits, batch_targets)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

print(batch_features.shape, batch_targets.shape)

The exact final batch contents depend on the shuffled order, but every training example is visited once in the epoch. The printed shapes are stable: with twenty examples and a batch size of five, each batch contains five rows, four input columns, and one target value per row.

Design Choices and Trade-Offs

batch_size is a memory and optimization decision. Larger batches improve hardware utilization and produce smoother gradient estimates, but they use more memory and may generalize differently. Smaller batches are noisier and often fit limited hardware, but they spend more time in Python overhead and optimizer steps.

Shuffling is usually correct for training but usually wrong for evaluation reports that must be reproducible by row order. For time series, naive shuffling can leak future context into earlier training windows if examples were constructed incorrectly. Shuffle completed training windows, not raw chronological records that still need causal slicing.

Preprocessing location is another trade-off. Cheap deterministic tensor conversions can live in __getitem__. Expensive resizing, tokenization, or augmentation may benefit from workers or offline preprocessing. Random augmentation inside workers should be deliberately seeded when experiments need repeatability.

Failure Modes and Troubleshooting

Symptom: RuntimeError says tensors have unequal sizes during collation. Cause: examples return tensors with different shapes and the default collator tries to stack them. Diagnose: run the dataset directly for several indices and print each field shape before using the loader. Correct: resize inputs, pad them in collate_fn, or return a list field that the model handles explicitly.

Symptom: labels no longer match inputs after shuffling. Cause: features and targets were stored in separate datasets or shuffled independently before construction. Diagnose: include a stable sample id in each item and inspect ids, features, and labels in the same batch. Correct: create one dataset object that returns the aligned pair from the same index.

Symptom: loading works with num_workers=0 but hangs or crashes with workers. Cause: the dataset contains unpicklable state, unsafe global handles, or code that behaves differently in child processes. Diagnose: reproduce with one worker, remove side effects from __init__, and test dataset[0] in isolation. Correct: open files lazily in workers, avoid shared mutable state, and keep worker count modest until measured.

Reliability and Performance Implications

Data pipelines influence model reliability directly. A loader that silently drops the last minority-class examples, applies random transforms to validation data, or changes label dtype can produce misleading metrics. Keep training, validation, and test loaders separate, and make their shuffle and augmentation policies visible in experiment logs.

Performance should be measured as examples per second and device utilization, not just wall-clock epoch time. If the GPU waits between batches, increase num_workers, precompute expensive transforms, use pinned memory for CUDA transfers, or simplify per-item Python work. If CPU memory grows, reduce prefetching, batch size, or cached objects held by each worker.

Hands-On Lab

Prerequisites: a Python environment with PyTorch installed and a terminal where you can run a single script. No external dataset is required.

Create a script that defines ten feature rows, ten labels, and ten string ids.
Implement a Dataset whose __getitem__ returns all three fields from the same index.
Create a training loader with batch_size=4 and shuffle=True.
Create a validation loader over the same toy dataset with batch_size=4 and shuffle=False.
Print the ids from both loaders for two epochs.
Verify that the training order changes between epochs while the validation order stays stable.
Change the dataset to return one variable-length tensor field, observe the default collation failure, then add a padding collate_fn.
Cleanup by deleting the script or reverting the toy variable-length change after saving the observed error message.

Successful verification means each epoch visits every id exactly once, labels stay attached to ids, validation order is deterministic, and variable-length examples batch only after the custom collator is installed.

Assessment Exercises

You have image tensors shaped [3, 224, 224] and scalar class labels. What batch shapes should a loader with batch_size=32 produce, and why?
A validation metric changes every run even though the model checkpoint is fixed. Which loader options and transforms would you inspect first?
A dataset reads one row from a CSV in __getitem__. Training is slow with low GPU utilization. Propose two improvements and explain the trade-off of each.
Why is it safer for one dataset item to return both the feature tensor and target tensor instead of using separate shuffled feature and label loaders?
Design a collate_fn output for token sequences that preserves enough information for a recurrent or attention model to ignore padding.

Summary

Dataset defines how to retrieve one example. DataLoader defines how examples become batches. Batching controls memory use and gradient behavior, shuffling controls training order, samplers express advanced index policies, and collation turns a list of examples into tensors the model can consume. Debug data pipelines from the item outward: inspect one example, inspect one batch, then inspect a full epoch before trusting a training curve.