Data loaders: preparing examples in parallel

A data loader turns a list of examples into a stream of batches

A data loader needs three things:

• A dataset: an object that knows how many examples there are and can prepare example number i on request.
• An order: usually a new random shuffle each epoch.
• A batch size: how many examples to stack into each batch. Stacking the examples into one array is called collating.

The snippet below builds all three in PyTorch. The dataset's two methods are the whole contract: __len__ says how many examples there are, and __getitem__ prepares one. The loader calls them for you. (Reading Python skips classes. The Python Tutorial explains the class syntax in A First Look at Classes.)

loader.py

import torch
# Load two names from torch.utils.data, so the code can write them without that prefix.
from torch.utils.data import Dataset, DataLoader

class Squares(Dataset):
    # Ten examples; example i is the pair (i, i*i).
    def __len__(self):
        return 10
    def __getitem__(self, i):              # prepare one example
        # A tensor, PyTorch's array, that holds i and i * i as 32-bit decimal numbers.
        return torch.tensor([i, i * i], dtype=torch.float32)

if __name__ == "__main__":
    # fix the random shuffle, so each run prints the same order
    torch.manual_seed(0)
    loader = DataLoader(Squares(), batch_size=4, shuffle=True, num_workers=2)
    for epoch in range(2):
        for batch in loader:
            # Print the batch's shape, then the i of each example in it as a list
            # of whole numbers.
            print(epoch, tuple(batch.shape), batch[:, 0].int().tolist())

0 (4, 2) [6, 7, 1, 4]
0 (4, 2) [2, 0, 9, 8]
0 (2, 2) [3, 5]
1 (4, 2) [2, 4, 7, 0]
1 (4, 2) [8, 9, 5, 3]
1 (2, 2) [6, 1]

Each epoch visits all ten examples once, in a new order. The last batch is short, because 10 is not a multiple of 4. num_workers=2 starts two worker processes. A process is a separate running program with its own memory, so each worker can run on its own CPU core at the same time as the others. Some systems start each worker as a fresh Python that reads the file again. The if __name__ == "__main__": line keeps such a worker from running the training code too.

Too few workers leave the GPU waiting

Each worker prepares a whole batch, then starts on another one. The training loop takes the batches in order. If the workers together prepare batches more slowly than the GPU uses them, the GPU waits.

Try it. The timeline below starts with one worker. Predict how many workers you need for the numbers in the worked example. Then move the slider until the blue training blocks touch with no gaps. Then make preparation slower and find the new number.