Overview

In our first PyTorch classifier, we passed the entire training partition to the model at once. That was convenient for a small example. What if we want to work through a larger dataset a few examples at a time?

PyTorch gives us two useful pieces: a Dataset that supplies examples and a DataLoader that gathers them into batches^†. Let's start with our clowns, then use these pieces with the classifier we've already seen.

Objectives

retrieve an example and its matching label from a Dataset
assemble batches with a DataLoader and predict their shapes
preserve feature-label pairing when shuffling examples
adapt our earlier training loop to work with batches

Getting started

In order to follow along, ensure you've configured your development environment with a working docker installation.

Introduction

You can execute the code snippets that follow in the provided docker image:

docker run -it -p 127.0.0.1:8888:9999 git.hlt.parsertongue.org/teaching/docker-images/pytorch:latest

Open your browser to localhost:8888, enter the token printed in the terminal, and create a new notebook (New $\rightarrow$ Python 3 (ipykernel)).

Let's import the pieces we'll use:

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

SEED = 42

Start with familiar tensors

Recall the clowns in Tensor Basics. Each row has height, nose diameter, and shoe length, all in centimeters. We'll keep those same three clowns and attach an illustrative label: whether their act includes juggling (1 for yes, 0 for no).

clowns = torch.tensor([
    [170, 6, 45],
    [185, 8, 50],
    [160, 7, 55],
], dtype=torch.float32)
juggles = torch.tensor([1, 0, 1], dtype=torch.float32)

These labels are made up. We're using them to inspect how features and labels travel together, not to suggest that nose diameter predicts juggling! We won't train a classifier on these three examples.

Example	Features (cm)	Juggles
Clown 0	`[170, 6, 45]`	1
Clown 1	`[185, 8, 50]`	0
Clown 2	`[160, 7, 55]`	1

Our feature tensor has shape (3, 3), and the labels have shape (3,). The first axis runs over the same examples in both tensors. We'll use floating point labels here because we'll later return to binary cross entropy with sigmoid probabilities.

Retrieve one example with a Dataset

For this kind of indexable dataset, PyTorch needs two methods: __len__ tells us how many examples there are, and __getitem__ retrieves one. The double underscores connect these 'dunder' methods to operations such as len(dataset) and dataset[1].

class ClownDataset(Dataset):
    def __init__(self, features, labels):
        if len(features) != len(labels):
            raise ValueError("Each example needs a matching label.")
        self.features = features
        self.labels = labels

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

    def __getitem__(self, index):
        # Use the same index for the features and their label.
        return self.features[index], self.labels[index]

dataset = ClownDataset(clowns, juggles)
features, label = dataset[1]

print(len(dataset))       # 3
print(features.tolist())  # [185.0, 8.0, 50.0]
print(label.item())       # 0.0

dataset[1] gives us one example, not a batch. Its feature vector has shape (3,), and its label is a scalar tensor with shape (). The dataset doesn't add a batch axis or decide the batch size.^†

Assemble a batch with a DataLoader

Now let's ask for two examples at a time.^† We'll leave shuffling off while we inspect the result:

loader = DataLoader(
    dataset,
    batch_size=2,
    shuffle=False,
    drop_last=False,
    num_workers=0,
)

for batch_x, batch_y in loader:
    print(batch_x.tolist(), batch_y.tolist())
    print(tuple(batch_x.shape), tuple(batch_y.shape))

The first batch contains clowns 0 and 1. The second contains clown 2:

Batch	Feature shape	Label shape	Labels
First	`(2, 3)`	`(2,)`	`[1.0, 0.0]`
Second	`(1, 3)`	`(1,)`	`[1.0]`

The loader retrieves individual examples, then stacks their feature vectors into rows and their labels into a vector. This step is called collation. The default collation adds a batch axis while preserving the order of features within each example. It doesn't transpose our features.

num_workers=0 keeps data loading in the current process, which is a simple starting point for this notebook. More workers can help with expensive data loading, but they aren't necessary to understand batching.

Shuffle examples, not their parts

For ordinary training, we usually want a new ordering of examples each epoch. A loader can do that for us:

shuffled_loader = DataLoader(
    dataset,
    batch_size=2,
    shuffle=True,
    generator=torch.Generator().manual_seed(SEED),
    num_workers=0,
)

With this loader, clown 2 might appear before clown 0. Clown 2's label still travels with clown 2's features. The loader shuffles example indices, not the columns within a feature vector and not the labels independently. Our explicit generator makes the sequence of shuffled orders reproducible in this example. It doesn't mean every epoch uses the same order.

Use `TensorDataset` when the tensors are ready

Our custom class only indexes tensors. TensorDataset already does that:

tensor_dataset = TensorDataset(clowns, juggles)
features, label = tensor_dataset[1]

assert torch.equal(features, clowns[1])
assert torch.equal(label, juggles[1])

The tensors must have the same length along their first axis. They don't need identical shapes: our features are (3, 3) and our labels are (3,). At each index, TensorDataset returns the corresponding item from each tensor. For this case, it's a useful alternative to writing our own class.

Return to our classifier

Let's leave the clowns in their car again and revisit the two-feature synthetic data from First steps with PyTorch. This block regenerates that data so the tutorial can run in its own notebook:

from sklearn.datasets import make_blobs
from sklearn.model_selection import train_test_split

X, y = make_blobs(
    n_samples=100,
    n_features=2,
    centers=2,
    random_state=SEED,
)
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=SEED,
)

# Keep each feature row paired with its target.
train_dataset = TensorDataset(
    torch.from_numpy(X_train).float(),
    torch.from_numpy(y_train).float(),
)
test_dataset = TensorDataset(
    torch.from_numpy(X_test).float(),
    torch.from_numpy(y_test).float(),
)

train_loader = DataLoader(
    train_dataset,
    batch_size=32,
    shuffle=True,
    generator=torch.Generator().manual_seed(SEED),
    num_workers=0,
)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False, num_workers=0)

There are 80 training examples, so each training pass has batches of 32, 32, and 16 examples. We keep the test loader's order unchanged to make its predictions easier to inspect.

We'll use the same linear layer and sigmoid as before, this time joined with nn.Sequential:

torch.manual_seed(SEED)
model = nn.Sequential(nn.Linear(2, 1), nn.Sigmoid())
loss_fn = nn.BCELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

The important change is the inner loop over batches. We calculate a loss and update the parameters after each batch, rather than after passing all 80 examples through the model:

EPOCHS = 5
loss_tracker = []
model.train()

for epoch in range(EPOCHS):
    epoch_loss = 0.0
    for batch_x, batch_y in train_loader:
        # Clear the previous batch's gradients before calculating new ones.
        optimizer.zero_grad()

        probabilities = model(batch_x).squeeze(-1)
        loss = loss_fn(probabilities, batch_y)
        loss.backward()
        optimizer.step()

        # Weight each batch mean by the number of examples it contains.
        epoch_loss += loss.item() * len(batch_x)

    mean_loss = epoch_loss / len(train_dataset)
    loss_tracker.append(mean_loss)
    print(f"Epoch {epoch + 1}: mean training loss {mean_loss:.3f}")

An epoch is one complete pass over the training data. Here, an epoch contains three optimizer updates, not just one. This example illustrates the loop, not a search for the best batch size or training duration. We'll revisit those choices in training tips.

Evaluation can use batches too. We switch to evaluation mode and turn off gradient recording, then collect the results in the test loader's order:

model.eval()
probability_batches = []
with torch.no_grad():
    for batch_x, batch_y in test_loader:
        probability_batches.append(model(batch_x).squeeze(-1))

test_probabilities = torch.cat(probability_batches)
test_predictions = (test_probabilities > 0.5).long()
print(tuple(test_predictions.shape))  # (20,)

The data loader supplies data. It doesn't calculate gradients, update weights, or switch the model into evaluation mode. Those are still our training and evaluation code's responsibilities.

Carry the same idea into the CNN

The small text CNN expects rows of character IDs. We can use the same batching pattern:

# Same fixed-length ID rows as the CNN example.
ids = torch.tensor([[1, 2, 3, 1, 2], [3, 1, 2, 3, 1]], dtype=torch.long)
class_ids = torch.tensor([0, 1], dtype=torch.long)
text_dataset = TensorDataset(ids, class_ids)
text_loader = DataLoader(text_dataset, batch_size=2, shuffle=False, num_workers=0)

batch_ids, batch_classes = next(iter(text_loader))
print(tuple(batch_ids.shape))      # (2, 5)
print(tuple(batch_classes.shape))  # (2,)

Here, the inputs stay integer IDs for the embedding lookup. The targets are integer class indices for a two-class classifier using CrossEntropyLoss, rather than the floating point binary targets we used with sigmoid and BCELoss. The CNN can turn (batch, positions) IDs into embeddings, then transpose the position and channel axes as before. Batching doesn't replace those model operations.

Practice

Customize data loading or write a custom DataLoader?

What if our data needs something the examples above don't cover? We can usually customize the existing loader without writing a new class. A custom Dataset changes how we retrieve an example. A collate_fn changes how examples become a batch, such as padding texts to the longest sequence in that batch. A sampler chooses which example indices to retrieve and their order, while a batch_sampler groups those indices into batches.

Writing a DataLoader subclass or a wrapper can make sense when a project's iteration behavior goes beyond these hooks. That's a more advanced case, not something we need just because our data is unusual. Start with the part of the pipeline you actually need to change. PyTorch's data loading documentation describes these customization points.

Before running a variation, predict what will change:

Set the clown loader's batch size to 1. What are the feature and label shapes now? Has each clown's feature count changed?
Set drop_last=True with batch size 2. Which clown is omitted when shuffling is off? What could happen to that omission across shuffled epochs?
Try a training batch size that leaves a final batch of one example. Why does squeeze(-1) still work?

The aim isn't to write a complicated data pipeline. Retrieve one example correctly, preserve its label, then check the batch axes. Those habits carry over when the examples become texts and the model becomes a CNN.

Datasets and DataLoaders in PyTorch

Overview

Objectives

Getting started

Introduction

Start with familiar tensors

Retrieve one example with a Dataset

Assemble a batch with a DataLoader

Shuffle examples, not their parts

Use `TensorDataset` when the tensors are ready

Return to our classifier

Carry the same idea into the CNN

Practice

Further reading

Overview

Objectives

Getting started

Introduction

Start with familiar tensors

Retrieve one example with a Dataset

Assemble a batch with a DataLoader

Shuffle examples, not their parts

Use TensorDataset when the tensors are ready

Return to our classifier

Carry the same idea into the CNN

Practice

Further reading

Use `TensorDataset` when the tensors are ready