Datasets and DataLoaders

The idea

The previous module sliced tensors by hand to make mini-batches. That works for 300 rows in memory. Real projects have data that doesn't fit in memory (thousands of images), needs per-example processing (reading a file, augmenting an image), and should load in parallel while the GPU trains. PyTorch splits this into two pieces:

• a Dataset: knows how many examples there are and how to fetch example number i;
• a DataLoader: draws examples from a Dataset, shuffles them, groups them into batches and (optionally) loads them in parallel.

A Dataset is two methods

In the last module, the training loop sliced batches out of tensors by hand. That works for small data, but it's fiddly and it doesn't scale to data that doesn't fit in memory. PyTorch's answer is to wrap the data in a Dataset: a class that knows how many examples there are and how to fetch one:

import numpy as np
import pandas as pd
import torch
from torch.utils.data import Dataset, DataLoader


class ChurnDataset(Dataset):
    def __init__(self, frame, features, target):
        self.X = torch.tensor(frame[features].to_numpy(), dtype=torch.float32)
        self.y = torch.tensor(frame[target].to_numpy(), dtype=torch.float32)

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

    def __getitem__(self, i):
        return self.X[i], self.y[i]


churn = pd.read_csv("churn.csv")
FEATURES = ["tenure_months", "orders_90d", "days_since_last_order", "support_tickets", "app_user"]
dataset = ChurnDataset(churn, FEATURES, "churned")
print(len(dataset))
x0, y0 = dataset[0]
print(x0, y0)

400
tensor([12.,  7., 47.,  1.,  0.]) tensor(0.)

• __len__ returns the number of examples.
• __getitem__(i) returns example i, usually as (features, target).

That's the whole contract (Python's special methods, from the classes module). For data already in tensors, TensorDataset(X, y) does the same in one line.

A DataLoader makes batches

A Dataset hands out one example at a time. Training wants shuffled batches of them. A DataLoader sits on top of the Dataset and does exactly that:

torch.manual_seed(0)
loader = DataLoader(dataset, batch_size=64, shuffle=True)
print(len(loader))
X_batch, y_batch = next(iter(loader))
print(X_batch.shape, y_batch.shape)
print(sum(len(yb) for _, yb in loader))

7
torch.Size([64, 5]) torch.Size([64])
400

• batch_size=64: each iteration yields 64 examples stacked into tensors ((64, 5) features, (64,) targets).
• shuffle=True: a new random order every epoch. Use it for training; not for validation and test, where order doesn't matter and reproducibility does.
• len(loader) is the number of batches per epoch (400 / 64 → 7, the last one smaller).
• num_workers=4 (not shown) loads batches in background processes: essential for images, unnecessary for small tables.

Splitting a Dataset

random_split divides a Dataset into training and validation parts:

from torch.utils.data import random_split

generator = torch.Generator().manual_seed(42)
train_set, val_set = random_split(dataset, [300, 100], generator=generator)
train_loader = DataLoader(train_set, batch_size=32, shuffle=True)
val_loader = DataLoader(val_set, batch_size=100)
print(len(train_set), len(val_set), len(train_loader), len(val_loader))

300 100 10 1

The seeded generator makes the split reproducible, like random_state in scikit-learn. (For stratified splits or preprocessing fitted on training rows only, splitting the DataFrame with scikit-learn first and building two Datasets is often simpler.)

The cleaner training loop

With loaders, the loop from the previous module becomes:

from torch import nn

mean = dataset.X[train_set.indices].mean(dim=0)
std = dataset.X[train_set.indices].std(dim=0)

torch.manual_seed(0)
model = nn.Sequential(nn.Linear(5, 16), nn.ReLU(), nn.Linear(16, 1))
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
loss_fn = nn.BCEWithLogitsLoss()

for epoch in range(1, 31):
    model.train()
    for X_b, y_b in train_loader:
        optimizer.zero_grad()
        loss = loss_fn(model((X_b - mean) / std).squeeze(1), y_b)
        loss.backward()
        optimizer.step()

    model.eval()
    correct = 0
    with torch.no_grad():
        for X_b, y_b in val_loader:
            probs = torch.sigmoid(model((X_b - mean) / std).squeeze(1))
            correct += ((probs >= 0.5).float() == y_b).sum().item()
    if epoch in (1, 10, 30):
        print(f"epoch {epoch:2d}  validation accuracy {correct / len(val_set):.2f}")

epoch  1  validation accuracy 0.73
epoch 10  validation accuracy 0.86
epoch 30  validation accuracy 0.84

The loop body no longer knows how batches are made; the loader handles it. The same loop would work unchanged for a dataset of millions of images on disk.

Note the details: squeeze(1) turns the model's (batch, 1) output into (batch,) to match the targets, and the scaling statistics come from the training rows only.

Why this matters for clean code

Datasets separate data access from training logic. You can test a Dataset on its own (does dataset[0] return the right shapes?), swap a CSV-backed Dataset for a database-backed one without touching the loop, and reuse the same loop across projects. It's the same separation as load_data() vs analyse() in the pandas pipelines module.

Recap

• Dataset: implement __len__ and __getitem__(i); TensorDataset(X, y) for in-memory tensors.
• DataLoader(dataset, batch_size, shuffle=True) yields batches; no shuffling for validation/test.
• len(loader) = batches per epoch; num_workers loads in parallel.
• random_split(dataset, [n_train, n_val], generator=...) for reproducible splits.
• Training loops iterate for X_b, y_b in loader:; fit preprocessing on training rows only.