D4.4 · LESSON 4 OF 6

Data and minibatches

Translate an inspected NumPy model into tensors, modules, automatic differentiation and a clear training loop.

THE IDEA

Feeding the model in batches

You rarely train on the whole dataset in one step. A DataLoader walks the dataset and hands you batches: tensors whose first axis is the batch axis. Each row of a batch is one independent example, and the same weights are applied to every row.

batch_size sets how many examples per step. If the dataset size is not a multiple, the last batch is smaller.
drop_last=True discards that smaller batch.
shuffle=True visits the examples in a new random order each epoch. Features and targets are shuffled together, so they stay paired.

batch_size4

shuffle=True drop_last=True

batch 1 (4)sample 0sample 1sample 2sample 3

batch 2 (4)sample 4sample 5sample 6sample 7

batch 3 (2)sample 8sample 9

dataset of 10 samples (indices 0..9) loader = DataLoader(ds, batch_size=4, shuffle=False, drop_last=False) len(loader) = 3 batch sizes: 4, 4, 2

The shuffled order here comes from a small generator of our own. Real PyTorch draws a different permutation, but the rules are the same: every sample once per epoch, X and y shuffled together.

Ten samples. Change the settings and see exactly which samples fall in which batch.

Predict first. What does the batch axis of a tensor represent?

real PyTorch (run it locally)

from torch.utils.data import DataLoader, TensorDataset
loader = DataLoader(TensorDataset(X, y), batch_size=4, shuffle=True, drop_last=False)
for xb, yb in loader:          # xb: (4, features), yb: (4,) ... the last batch may be smaller
    ...

Chapter context

PyTorch provides tensor operations and differentiation infrastructure. The framework removes bookkeeping, but it does not remove the need to understand dimensions, parameter ownership, device placement and evaluation behavior.

TRY IT YOURSELF

Part 1. Do it hands-on

The loader lab below is a visual model of a real DataLoader over 10 samples. The Python part has you write the batching yourself with NumPy.

batch_size1

shuffle=True drop_last=True

batch 1 (1)sample 0

batch 2 (1)sample 1

batch 3 (1)sample 2

batch 4 (1)sample 3

batch 5 (1)sample 4

batch 6 (1)sample 5

batch 7 (1)sample 6

batch 8 (1)sample 7

batch 9 (1)sample 8

batch 10 (1)sample 9

dataset of 10 samples (indices 0..9) loader = DataLoader(ds, batch_size=1, shuffle=False, drop_last=False) len(loader) = 10 batch sizes: 1, 1, 1, 1, 1, 1, 1, 1, 1, 1

The shuffled order here comes from a small generator of our own. Real PyTorch draws a different permutation, but the rules are the same: every sample once per epoch, X and y shuffled together.

To do: 10 samples, in order: set the loader so the batches have sizes 4, 4 and 2
To do: Keep only full batches: sizes 4 and 4 (drop_last=True)
To do: Find a batch size without drop_last where the last batch contains a single sample
To do: Turn shuffling on (batch size 4) and pick a seed for which sample 0 is NOT in the first batch

Part 2. Do it in code

batch_indices(n, batch_size, drop_last=False): return a list of lists with the indices 0..n-1 split in order into batches. Without drop_last the last batch may be shorter; with drop_last=True a short last batch is discarded. batch_size < 1 raises ValueError.
make_batches(X, y, batch_size, shuffle=False, seed=0, drop_last=False): X is an (n, d) numpy array and y an (n,) array. Return a list of (Xb, yb) pairs. If shuffle is True, use order = np.random.default_rng(seed).permutation(n); otherwise order = np.arange(n). The SAME order must be applied to X and y. Raise ValueError if the lengths differ.

import numpy as np

def batch_indices(n, batch_size, drop_last=False):
    pass

def make_batches(X, y, batch_size, shuffle=False, seed=0, drop_last=False):
    pass

Ctrl/⌘ + Enter runs it

Stuck, or something not working? You can come back to it any time.

CHECK YOUR UNDERSTANDING

What does a batch axis represent?

AThe time steps of a sequence, always processed one after another.BIndependent examples processed together: each row of the batch is one example, and the same weights are applied to every row.CThe features of one example, because the batch axis lists the input features.

Chapter reference: math, worked example & code

This section belongs to the whole chapter, so it is the same on every lesson in it and may cover a different lesson's topic. The code for the experiment you just ran is inside the experiment itself.

Math and notation

For a batch X of shape (B,d), a dense output with h features has shape (B,h). torch.nn.Linear stores weight with shape (h,d) and computes X times weight-transpose plus bias. Parameter count is d*h+h. PyTorch execution belongs in a linked notebook or supported runtime; the chapter does not assume browser PyTorch.

Worked example

A batch has shape (2,4). A Linear(4,3) layer stores a (3,4) weight and a length-three bias. Its output is (2,3). The parameter count is 4×3+3=15. This differs from the batch's element count: model parameters are shared across both examples.

Reference code

The following example uses the Python standard library and does not call a remote service.

batch, inputs, outputs = 2,4,3
weight_shape = (outputs,inputs)
output_shape = (batch,outputs)
assert weight_shape == (3,4)
assert output_shape == (2,3)
assert inputs*outputs+outputs == 15

Python examples are provided for study and local execution; this page does not execute Python.

Common mistakes

eval mode and no_grad are independent. Model parameters do not multiply with batch size. A tensor detached for logging must not replace the tensor used to compute the training loss.