Datasets and Batches

Imagine trying to learn from a million photographs by looking at all of them at once, before making a single correction. It would be slow and wasteful. Instead, networks learn from small handfuls at a time. They take a quick look, make a small correction, and move on. That's why deep learning can scale.

• A million photos at once?
• Learn from small handfuls
• Quick corrections, many times

Dataset and DataLoader

PyTorch splits this job in two. A dataset knows two things. How many items there are, and how to get item number i. Here's a tiny one, where the target is the square of the input. A data loader then wraps it, and handles the rest. It shuffles the data, groups it into batches, and even loads in parallel. Run it. Each batch holds four shuffled items. If the last one is shorter, that's fine. These two classes work the same way for images, text, and audio.

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

class Squares(Dataset):
    def __init__(self, n):
        self.x = torch.arange(n, dtype=torch.float32)

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

    def __getitem__(self, i):
        return self.x[i], self.x[i] ** 2

loader = DataLoader(Squares(10), batch_size=4, shuffle=True)

for x, y in loader:
    print("x:", x.tolist(), " y:", y.tolist())

x: [4.0, 7.0, 1.0, 6.0]  y: [16.0, 49.0, 1.0, 36.0]
x: [2.0, 0.0, 8.0, 9.0]  y: [4.0, 0.0, 64.0, 81.0]
x: [5.0, 3.0]  y: [25.0, 9.0]

Why batches work so well

Batches are a lovely compromise. One example at a time is noisy, and it makes poor use of the G P U. The whole dataset at once is accurate, but very slow, and it may not even fit in memory. A batch of around thirty two to two hundred and fifty six gets the best of both. And there's a bonus. A little noise in the gradients helps the network escape from bad spots. One full pass over the data is called an epoch.

Summary