5.3.1 Dataset and DataLoader

Why Dataset and DataLoader?

In Gradient Descent, we trained with mini-batches by hand: shuffle the row numbers at the start of each epoch, cut them into groups of 16, and take one step per group. In Multi-Layer Perceptrons, we skipped batches altogether and pushed all 300 points through the network at once. Neither is how real PyTorch code looks. An IOAI task can have tens of thousands of rows or images, far too many for one step, and nobody wants to rewrite the shuffling and slicing for every task.

PyTorch splits the job into two small pieces. A Dataset knows how many samples there are and how to fetch sample number i. A DataLoader takes a Dataset and hands out shuffled mini-batches, one per training step. Every PyTorch training script you will read uses this pair, and the same two pieces carry over unchanged to images, text and audio later in the guide.

What you need to know

By the end of this module, you should be comfortable with these ideas:

Learn it properly

This page shows every idea with a small example that you can run. The course below explains the same pieces from a different angle, which helps them stick.

Pick one main course

The topic is small, so one short resource is enough.

Our pick: PyTorch's own tutorial

Read Datasets & DataLoaders from PyTorch's Learn the Basics series. It loads a dataset of clothing images, writes a custom Dataset for image files listed in a CSV, and loops over a DataLoader. The first run downloads the images, so do it at home rather than in a contest. The images themselves do not matter yet: focus on the three methods of the custom Dataset and on what comes out of the DataLoader.

If you prefer something else

• A video: Patrick Loeber's Dataset and DataLoader, Batch Training builds a Dataset from a CSV file of wine measurements, which is very close to a tabular IOAI task.
• A textbook: section 3.3 of Dive into Deep Learning, Synthetic Regression Data, writes a data loader by hand and then replaces it with PyTorch's classes. Seeing both makes it clear what the DataLoader does.
• The reference: the torch.utils.data documentation lists every option. You do not need to read it now, but it is the place to look something up.

A quick tour

The rest of this page builds a Dataset and a DataLoader for a tiny table, one piece at a time. Go through your chosen course first, then run each cell in a notebook as you read.

The data

We will use ten made-up students, with the hours each one studied and slept before an exam, and whether they passed. You can download students.csv.

A model takes tensors, so the first job is to turn the two feature columns and the label column into float32 tensors:

import pandas as pd
import torch

students = pd.read_csv("students.csv")
feature_columns = ["hours_studied", "hours_slept"]

features = torch.tensor(students[feature_columns].values,
                        dtype=torch.float32)
labels = torch.tensor(students["passed"].values, dtype=torch.float32)
features.shape, labels.shape

(torch.Size([10, 2]), torch.Size([10]))

Ten rows with two features each, and ten labels. The dtype=torch.float32 matters: without it, whole numbers become int64 tensors, which a model refuses (see Common mistakes).

A Dataset: a length and a way to get one sample

When the data is already in tensors, PyTorch's TensorDataset turns it into a Dataset in one line:

from torch.utils.data import TensorDataset

dataset = TensorDataset(features, labels)
len(dataset)

10

Asking for sample 0 gives that student's features and label, as a pair:

dataset[0]

(tensor([1., 6.]), tensor(0.))

That is the whole idea of a Dataset: len(dataset) says how many samples there are, and dataset[i] returns sample i. Nothing more is needed.

Writing your own Dataset

TensorDataset is enough for a table that fits in memory. When each sample needs work of its own, such as opening an image file or cleaning a sentence, you write a class instead. It is the ScoreDataset pattern from Classes and Inheritance: inherit from Dataset and fill in three methods. Here is one for our table:

from torch.utils.data import Dataset

class StudentDataset(Dataset):
    def __init__(self, table):
        columns = ["hours_studied", "hours_slept"]
        self.features = torch.tensor(table[columns].values,
                                     dtype=torch.float32)
        self.labels = torch.tensor(table["passed"].values,
                                   dtype=torch.float32)

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

    def __getitem__(self, index):
        return self.features[index], self.labels[index]

dataset = StudentDataset(students)
len(dataset), dataset[3]

(10, (tensor([3., 8.]), tensor(1.)))

Let us read it method by method:

• __init__ runs once. It does the slow work up front: here, converting the whole table to tensors.
• __len__ returns the number of samples.
• __getitem__ returns one sample. The DataLoader calls it many times per epoch, so keep it quick.

For images, the same three methods look slightly different: __init__ stores the list of file names, and __getitem__ opens one file and turns it into a tensor. You will write exactly that in Loading Images into Tensors with Torchvision.

A DataLoader makes the batches

A DataLoader wraps the Dataset and hands out batches when you loop over it. We start without shuffling, so the batches are easy to follow:

from torch.utils.data import DataLoader

loader = DataLoader(dataset, batch_size=4, shuffle=False)
for x_batch, y_batch in loader:
    print(x_batch.shape, y_batch.shape)

torch.Size([4, 2]) torch.Size([4])
torch.Size([4, 2]) torch.Size([4])
torch.Size([2, 2]) torch.Size([2])

Ten samples in batches of four gives three batches: 4, 4, and the 2 that are left. Each batch is two tensors, because each sample was a pair. The DataLoader called dataset[i] for four samples and stacked their features into one tensor of shape (4, 2), and their labels into one of shape (4,). The new first dimension is the batch dimension, which is exactly what a model expects. Here is the first batch in full:

x_batch, y_batch = next(iter(loader))
print(x_batch)
print(y_batch)

tensor([[1., 6.],
        [2., 8.],
        [3., 5.],
        [3., 8.]])
tensor([0., 0., 0., 1.])

These are rows 0 to 3 of the table. next(iter(loader)) is a handy way to grab one batch and check its shapes before you train anything.

Shuffling

For training, we want a new random order every epoch, for the reason explained in Gradient Descent: sorted data would give batches that all pull the same way. To see the order the DataLoader picks, we can give it a plain range(10). A range has a length and can be indexed, so it counts as a Dataset, and each batch then shows which sample numbers it holds:

torch.manual_seed(0)
order_loader = DataLoader(range(10), batch_size=4, shuffle=True)

for epoch in range(2):
    print("epoch", epoch + 1)
    for batch in order_loader:
        print(batch)

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

Each epoch uses every sample exactly once, in a different order. torch.manual_seed(0) makes the order the same every time you run the cell, which is useful while you debug. Step through the same two epochs, with our ten students:

Batches in a training loop

A training loop with a DataLoader has two loops: one over epochs, and inside it one over batches. The five steps from PyTorch Basics now run once per batch. Here we train the simplest possible classifier, one nn.Linear layer, to predict passed:

from torch import nn

torch.manual_seed(0)
train_loader = DataLoader(dataset, batch_size=4, shuffle=True)
model = nn.Linear(2, 1)
loss_fn = nn.BCEWithLogitsLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

for epoch in range(1, 201):
    total_loss = 0.0
    for x_batch, y_batch in train_loader:
        scores = model(x_batch).squeeze(1)    # (4, 1) -> (4,)
        loss = loss_fn(scores, y_batch)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        total_loss += loss.item() * len(y_batch)
    if epoch == 1 or epoch % 50 == 0:
        print(epoch, round(total_loss / len(dataset), 3))

1 0.987
50 0.571
100 0.561
150 0.555
200 0.539

The loss goes down, so the loop works. The model itself is not the point here; A Simple MLP in PyTorch trains a real network with this same loop. Two details are worth noticing:

• BCEWithLogitsLoss is the loss for a yes/no label, and it wants the scores and labels to have the same shape, hence .squeeze(1). Loss Functions explains it.
• loss.item() is the average loss over one batch. The last batch is smaller, so to get a fair average over the epoch we multiply each batch's loss by its size, add them up, and divide by the number of samples.

Training and validation data

A model must be checked on data it did not train on. PyTorch can split a Dataset directly with random_split. We give it the sizes of the two parts and a seeded random generator, so the split is the same on every run:

from torch.utils.data import random_split

generator = torch.Generator().manual_seed(0)
train_part, val_part = random_split(dataset, [8, 2],
                                    generator=generator)
len(train_part), len(val_part), val_part.indices

(8, 2, [6, 2])

Each part is a Dataset too, so each gets its own DataLoader:

train_loader = DataLoader(train_part, batch_size=4, shuffle=True)
val_loader = DataLoader(val_part, batch_size=64, shuffle=False)
len(train_loader), len(val_loader)

(2, 1)

The validation loader does not shuffle, because the order does not change a score and a fixed order makes results easier to compare. It can also use bigger batches, because no gradients are stored during validation.

In real tasks, you will often split the table instead, with scikit-learn's train_test_split, and then build one Dataset from each part. That is the better route whenever you also scale the features: the scaler must learn its means and standard deviations from the training rows only, as Approaching a Typical Problem explains. PyTorch Pipeline for a Simple Machine Learning Problem does it that way from start to finish.

Settings you can ignore for now

The DataLoader has many more options. You will see these in other people's code, and none of them is needed for tabular data:

• num_workers loads batches in background processes. It can speed up image loading; for a table in memory, the default 0 is fine.
• pin_memory=True can make copying batches to the GPU slightly faster.
• collate_fn changes how samples are stacked into a batch. You need it when samples have different sizes, such as sentences of different lengths; Tokenization shows the simpler alternative, padding every sentence to the same length first.
• sampler changes which samples are picked, for example to show rare classes more often.

Common mistakes

The messages are exactly what PyTorch prints.

How to practise

There is no PyTorch section on the practice sites we use (CodeChef, CS50 and HackerRank), so practise with the course material and small experiments of your own:

• Run the tutorial's custom Dataset with a CSV file of your own, such as the Titanic data from the Titanic Walkthrough.
• Take the hand-written mini-batch loop from Gradient Descent and rewrite it with a TensorDataset and a DataLoader.
• Change batch_size in the training loop above to 1, 2 and 10, and predict how many steps each epoch takes before you run it.

Check yourself

What comes next

• A Simple MLP in PyTorch: a real network, trained with the loaders from this page.
• PyTorch Pipeline for a Simple Machine Learning Problem: the whole route from a CSV file to a submission file.
• Loading Images into Tensors with Torchvision: the same Dataset and DataLoader, for images.