The data loader

Up to here the whole dataset sat in a single tensor. That works with 64 rows and stops working at 64,000: the forward pass would have to hold all intermediate results in memory at once, and the backward pass once more on top.

The solution is banal and carries the same name in every framework: batch. Take 32 rows, do a complete step with them, take the next 32. One pass through all the chunks is an epoch. A side benefit falls out that nobody planned: many small steps on changing excerpts often learn better than a few big ones on everything - the noise helps.

The goal

Training is half data plumbing. The insight of this chapter is less about mathematics than about the job: the part that shows up in talks is the gradients. The part that eats the week is this one - shapes, orders, indices, the last incomplete batch.

A data loader has to do three things, and all three are index work:

Cutting. n rows turn into ceil(n / batch_size) blocks. The last one may be smaller - dropping it would be the worse choice, because then with an unlucky size the same examples are systematically missing.

Shuffling. Freshly before every epoch. Otherwise the network sees the examples in the same order every time, and if the data arrives sorted - first all the zeros, then all the ones - a batch consists of nothing but one class. The network then alternately learns "everything is zero" and "everything is one".

Keeping pairs. x and y have to be shuffled with the same permutation. Two separate shuffles produce examples with someone else's answers

• a bug that does not crash, it just makes sure nothing learns.

Zero magic in this step

NumPy is allowed, torch, jax, tinygrad and autograd are forbidden - and this time torch.utils.data explicitly as well. NumPy may do more here than elsewhere: shuffling indices is pure data administration and happens before the graph, not inside it. Only what is computed with x and y has to run through tensor operations.

The exercise

Two files. First meintorch/data.py:

import numpy as np

from .tensor import Tensor


class DataLoader:
    """Hands out (x, y) batches as tensors. One iteration is one epoch."""

    def __init__(self, x, y, batch_size=32, shuffle=True, seed=None):
        self.x = np.asarray(x, dtype=np.float64)
        self.y = np.asarray(y, dtype=np.float64)
        self.batch_size = batch_size
        self.shuffle = shuffle
        self.rng = np.random.default_rng(seed)

    def __len__(self):
        # TODO(du): number of batches. The last one may be smaller -
        # it still counts.
        raise NotImplementedError

    def __iter__(self):
        # TODO(du): build indices, shuffle them if asked, cut them into
        # blocks of batch_size and yield (Tensor(x), Tensor(y)) per block.
        # x and y with the same indices.
        raise NotImplementedError

Then a second training function in meintorch/train.py that runs over batches instead of over the whole dataset:

def fit(model, loader, epochs=10, lr=0.05, loss_fn=mse_loss):
    """Like train(), but one pass through all batches per epoch."""
    opt = SGD(model.parameters(), lr=lr)
    history = []
    for _ in range(epochs):
        # TODO(du): iterate over the loader, per batch the same five lines
        # as in chapter 06, average the loss values of one epoch and append
        # that average.
        raise NotImplementedError
    return history

Hint 1 - the direction

__iter__ is a generator: yield instead of return, and then for xb, yb in loader works on its own.

The trick when cutting is not to shuffle the data but the indices. An array 0..n-1, shuffled once, then cut into slices - and each slice is a selection index for x and for y alike. That way the pairing comes for free instead of having to be guarded.

For fit: the inner block is literally the one from chapter 06, only with xb, yb instead of x, y. The only difference is the extra loop around it.

Hint 2 - more concrete

• __len__: int(np.ceil(len(self.x) / self.batch_size)).
• Indices: idx = np.arange(len(self.x)), then self.rng.shuffle(idx) if self.shuffle is set. shuffle works in place.
• Cutting: for start in range(0, len(idx), self.batch_size), inside it part = idx[start : start + self.batch_size].
• Yielding: yield Tensor(self.x[part]), Tensor(self.y[part]).
• In fit: total += float(loss.data) per batch, at the end of the epoch history.append(total / len(loader)).

    # in data.py
    def __len__(self):
        return int(np.ceil(len(self.x) / self.batch_size))

    def __iter__(self):
        idx = np.arange(len(self.x))
        if self.shuffle:
            self.rng.shuffle(idx)
        for start in range(0, len(idx), self.batch_size):
            part = idx[start : start + self.batch_size]
            yield Tensor(self.x[part]), Tensor(self.y[part])


# in train.py
def fit(model, loader, epochs=10, lr=0.05, loss_fn=mse_loss):
    opt = SGD(model.parameters(), lr=lr)
    history = []
    for _ in range(epochs):
        total = 0.0
        for xb, yb in loader:
            loss = loss_fn(model(xb), yb)
            opt.zero_grad()
            loss.backward()
            opt.step()
            total += float(loss.data)
        history.append(total / len(loader))
    return history

In goal mode

The autonomous run has to prove the coverage, not just check the shapes:

Create `meintorch/data.py` with a class
`DataLoader(x, y, batch_size=32, shuffle=True, seed=None)`. `__iter__` shuffles
the indices (not the data), cuts them into blocks and yields
`(Tensor(x_batch), Tensor(y_batch))` per block; x and y are selected with the
same permutation. `__len__` returns the number of batches, the last one may be
smaller. Extend `meintorch/train.py` with
`fit(model, loader, epochs=10, lr=0.05, loss_fn=mse_loss)`, which runs over all
batches per epoch and returns the mean loss per epoch.
Mandatory: NumPy allowed, torch/jax/tinygrad/autograd forbidden, no
torch.utils.data either.
Write `tests/test_07_dataloader.py`: the batches of one epoch have to cover the
dataset exactly once, with 10 rows and batch_size=3 the sizes have to be
[3, 3, 3, 1], x and y have to stay paired, shuffling has to change the order and
not the content, and `fit` has to lower the loss.
Run `python -m pytest` until everything is green.

The milestone test

tests/test_07_dataloader.py. The milestone is the first test: all batches together add up to exactly the dataset - no example twice, none forgotten. The test data is deliberately numbered so that every row stays recognisable by its value.

import numpy as np

from meintorch.data import DataLoader
from meintorch.nn import Linear, ReLU, Sequential
from meintorch.train import fit


def data(n=10):
    # row i is [2i, 2i+1] with answer i - every row is recognisable
    x = np.arange(n * 2, dtype=float).reshape(n, 2)
    y = np.arange(n, dtype=float).reshape(n, 1)
    return x, y


def test_the_batches_cover_the_dataset_exactly_once():
    x, y = data(10)
    seen = np.concatenate(
        [xb.data for xb, _ in DataLoader(x, y, batch_size=3, seed=0)]
    )
    assert seen.shape == x.shape
    assert np.allclose(np.sort(seen[:, 0]), np.sort(x[:, 0]))


def test_the_last_batch_may_be_smaller():
    x, y = data(10)
    loader = DataLoader(x, y, batch_size=3, shuffle=False)
    assert [xb.shape[0] for xb, _ in loader] == [3, 3, 3, 1]
    assert len(loader) == 4


def test_x_and_y_stay_paired():
    x, y = data(10)
    for xb, yb in DataLoader(x, y, batch_size=4, seed=1):
        assert np.allclose(xb.data[:, 0] / 2.0, yb.data[:, 0])


def test_shuffling_changes_the_order_not_the_content():
    x, y = data(20)
    shuffled = np.concatenate(
        [xb.data for xb, _ in DataLoader(x, y, batch_size=5, seed=1)]
    )
    ordered = np.concatenate(
        [xb.data for xb, _ in DataLoader(x, y, batch_size=5, shuffle=False)]
    )
    assert not np.allclose(shuffled, ordered)
    assert np.allclose(np.sort(shuffled[:, 0]), np.sort(ordered[:, 0]))


def test_fit_runs_over_the_batches_and_lowers_the_loss():
    rng = np.random.default_rng(0)
    x = rng.normal(size=(64, 2))
    y = 2.0 * x[:, :1] - x[:, 1:] + 0.5
    net = Sequential(Linear(2, 8, seed=0), ReLU(), Linear(8, 1, seed=1))
    history = fit(net, DataLoader(x, y, batch_size=16, seed=0), epochs=30, lr=0.05)
    assert len(history) == 30
    assert history[-1] < history[0] / 5