Datasets And Dataloaders in Pytorch

Section OneWhy batching, and why a library

Every training loop so far has pushed the whole dataset through the model in one forward pass, computed one loss, and taken one step. That is batch gradient descent. It is correct, and for any dataset of real size it is unusable.

• Memory. A forward pass holds the inputs, every intermediate activation, and the autograd graph — peak memory scales with how many samples are in flight. 100,000 images at 224×224×3 is ~15 GB before a single activation. The result is an out-of-memory crash, not slow training.
• Convergence. Full-batch buys one update per epoch. The same 25 epochs in batches of 32 gives hundreds of updates for the same total computation — and each batch gradient is noisy but unbiased, which helps the model escape shallow minima and saddle points.

Slicing by hand gets you most of the way

batch_size = 2
n_samples  = len(X_train_tensor)

for epoch in range(epochs):
    # Simply walk the dataset in chunks of `batch_size`
    for start_idx in range(0, n_samples, batch_size):
        end_idx = start_idx + batch_size
        X_batch = X_train_tensor[start_idx:end_idx]
        y_batch = y_train_tensor[start_idx:end_idx]

        y_pred = model(X_batch)
        loss   = loss_function(y_pred, y_batch.view(-1, 1))

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

The maths here is right and the memory is bounded, so the objection to this loop is architectural, not mathematical. Four things are missing: a standard interface (this assumes an in-memory, sliceable tensor — swap in images on disk and every line changes), a home for transformations, shuffling (real files are usually sorted by class), and parallelism (the GPU idles while the CPU decodes).

Section TwoTwo abstractions, one division of labour

PyTorch supplies two classes in torch.utils.data. They are deliberately narrow, and the discipline of keeping them narrow is the whole design.

Section ThreeWriting a Dataset: the three methods

A custom dataset subclasses torch.utils.data.Dataset and overrides three methods. There is no fourth. Everything else PyTorch gives you for free.

from torch.utils.data import Dataset, DataLoader

class CustomDataset(Dataset):

    def __init__(self, features, labels):
        self.features = features
        self.labels   = labels

    def __len__(self):
        return self.features.shape[0]                    # how many samples

    def __getitem__(self, index):
        return self.features[index], self.labels[index]  # one (features, label) pair

__init__ — what you store

• Stores what is needed to locate a sample — tensors, a dataframe, file paths, a database connection. Do whole-collection conversions here (cast, reshape): once beats ten thousand times.
• It need not load the data. Store filenames and read each file only when requested — that is what keeps memory bounded.

__len__ — the contract

• Returns one integer: the number of samples.
• The loader derives the legal index range and the batch count from it. If it lies, the loader either over-runs the data or silently ignores part of it.

__getitem__ — one sample, never a batch

• Takes one index, returns one sample — conventionally (features, label).
• The batch dimension is added later, by collate. Return pre-batched data and it gets created twice — shapes silently become nonsense, with no error raised.

The dataset used throughout these notes

Ten samples, two features each, two balanced classes, generated once with a fixed random seed. Small enough to print in full, which means every claim below can be checked against this table rather than taken on trust.

toy_dataset = CustomDataset(X_toy, y_toy)
print('len:', len(toy_dataset), '| sample[2]:', toy_dataset[2])

• len(toy_dataset) works because __len__ is defined — Python's built-in len dispatches to it.
• toy_dataset[2] works because __getitem__ is defined — square-bracket indexing dispatches to it.
• These are ordinary Python protocols, not special PyTorch powers. That is exactly why the interface feels natural.

Section FourTransformations: where they belong

Raw data is rarely in the form a model wants. The question is not whether to transform, but where that code lives — and the answer is inside __getitem__. Three things follow, none available if you transform the whole array up front:

• Lazy, per sample. Only the current batch is transformed, so memory stays bounded even when the transformed form is far larger than the raw one.
• Different every epoch. A random crop or flip gives a fresh version of sample 3 each epoch — exactly what augmentation needs. Transform up front and every epoch sees the identical image.
• Runs inside the workers. The cost is parallelised and hidden behind training; work done in __init__ happens once, serially, in the main process.

Section FiveWiring the DataLoader

With a Dataset in hand, batching requires one line and no batching code of your own.

toy_loader = DataLoader(toy_dataset, batch_size=2, shuffle=False)

for batch_features, batch_labels in toy_loader:
    print(batch_features.tolist(), '->', batch_labels.tolist())

With shuffle=False the batches are indices [0,1], [2,3], [4,5], [6,7], [8,9] — check them against the table in Section 3 and every value matches. Ten samples divided into batches of two gives five iterations of the loop, and len(toy_loader) returns 5.

Where the batch dimension comes from

• __getitem__ returns features of shape [2] — two features, no batch axis. Collate stacks two of those into [2, 2], and the scalar labels (shape []) into [2].
• So the batch dimension is created by stacking. It exists nowhere in the Dataset — and every shape bug is a confusion about which side of that line you are on.

Section SixShuffling and samplers

Shuffling looks like a convenience flag. It is a correctness requirement, and the flag is thin shorthand over a real object.

Why order matters

• Real files arrive sorted — by class, date, patient, source.
• Walk that order and consecutive batches are highly correlated: the model adjusts hard toward one class, then undoes it for the next, so the gradient signal oscillates instead of accumulating.
• Batch statistics — the mean and variance batch normalisation relies on — become unrepresentative of the dataset.

Our toy set is balanced 5 and 5, but not sorted. Here is what it would look like if it were — sorted by label, batches of two, no shuffling:

The sampler underneath the flag

The shuffle argument does not implement shuffling. It selects a sampler — an object whose entire job is to produce the sequence of indices the loader will consume:

• shuffle=True and sampler= are mutually exclusive — passing both is an error, since they specify the same thing.
• Shuffle the training loader only. Evaluation uses shuffle=False: order does not affect the metric, but a stable order makes results reproducible.

Section Sevencollate_fn: how a list becomes a batch

After calling __getitem__ for every index in the chunk, the loader holds a plain Python list. collate_fn turns it into tensors: it stacks position-wise — all features into one tensor, all labels into another — adding a leading batch dimension to each. Traced on the first toy batch:

# What collate_fn receives — a list of two samples, each a tuple:
[ (tensor([ 1.0683, -0.9701]), tensor(1)),
  (tensor([-1.1402, -0.8388]), tensor(0)) ]

# What collate_fn returns — one tuple of two stacked tensors:
( tensor([[ 1.0683, -0.9701],
          [-1.1402, -0.8388]]),   # shape [2, 2]  =  [batch, features]
  tensor([1, 0]) )                # shape [2]     =  [batch]

Purely structural — a list of k tuples becomes a tuple of tensors with a leading dimension of k: shapes [2] and [] became [2, 2] and [2]. This is the moment the batch dimension is born.

Section EightParallelism: num_workers and the pipeline effect

Correctness is settled; this is about throughput — and one common waste: an expensive accelerator idling while one CPU thread decodes a JPEG. Take 10,000 samples, batch_size=32, num_workers=4 → 313 batches per epoch.

Choosing the values

• num_workers=0 (the default) loads inline in the main process — every millisecond decoding is a millisecond the GPU sits idle. Free for in-memory tensors; often the dominant cost for image and audio.
• More is not uniformly better. Each worker is a real process with memory overhead, and every batch must be serialised back. Once preparation is hidden behind computation, extra workers add overhead and nothing else.
• pin_memory=True speeds the host→GPU copy. Only meaningful on CUDA — on a CPU-only machine, leave it False.

Section NineParameter reference

Every argument you will realistically use, and the question each one answers.

Section TenThe training pipeline, restructured

Introducing the loader changes the shape of the training code itself. What was one loop becomes two, nested.

for epoch in range(epochs):
    y_pred = model(X_train_tensor)               # the entire dataset
    loss   = loss_function(y_pred, y_train_tensor.view(-1, 1))

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()                             # 1 update per epoch

for epoch in range(epochs):
    model.train()                                # training mode
    epoch_loss = 0.0

    for batch_features, batch_labels in train_loader:    # INNER: one iteration per batch

        y_pred = model(batch_features)                   # 1. forward
        loss   = loss_function(y_pred, batch_labels.view(-1, 1))   # 2. loss
        epoch_loss += loss.item()

        optimizer.zero_grad()                            # 3. clear old gradients
        loss.backward()                                  # 4. backpropagate
        optimizer.step()                                 # 5. update — once per BATCH

    print(f'Epoch {epoch + 1} | loss = {epoch_loss / len(train_loader):.4f}')

• The inner body runs len(train_loader) times per epoch, each a full forward–backward–update cycle.
• Updates per epoch go from one to the number of batches — at no extra total computation. The same data is processed either way.
• Divide by len(train_loader) when printing. Bare loss.item() reports the last batch only — noisy and not comparable across runs.

Push that batch size to its extremes and the three named regimes fall out of one mechanism:

• batch_size=10 (the whole set) → batch gradient descent, one update per epoch.
• batch_size=1 → stochastic gradient descent, ten noisy updates per epoch.
• Anything between → mini-batch gradient descent.
• Not three algorithms. One algorithm, one knob.

Section ElevenFinal assessment

Three questions, drawn at random from across the document. Answer them, submit for a score and per-question explanations, then fetch another three — the draw differs each time, so it is worth repeating.

Topic 10 · Deep Learning — Dataset & DataLoader · Interactive study notes