PyTorch The DataLoader
The DataLoader serves batches from a Dataset. Its batch_size setting decides how many rows each batch holds.
Build a loader
The loader wraps the dataset and a batch size. Five rows in batches of 2 make 3 batches.
loader = DataLoader(data, batch_size=2)
print(len(loader))Read one batch
One step of the loader gives one batch. The batch stacks 2 rows into one tensor.
xb, yb = next(iter(loader))
print(xb.shape)
print(yb.shape)Batch size sets the shapes
A batch size of 5 serves all rows at once. The model sees one batch of shape (5, 1).
wide = DataLoader(data, batch_size=5)
xb, yb = next(iter(wide))
print(xb.shape)
print(yb.shape)Example
import torch
from torch.utils.data import TensorDataset, DataLoader
features = torch.tensor([[2.0], [3.0], [4.0], [5.0], [6.0]])
targets = torch.tensor([55.0, 62.0, 71.0, 79.0, 88.0])
data = TensorDataset(features, targets)
loader = DataLoader(data, batch_size=2)
for i, (xb, yb) in enumerate(loader, start=1):
print(f"Batch {i}: inputs {xb.shape}, targets {yb.shape}")Key points
- DataLoader: an object that serves a Dataset in batches
- batch_size: how many rows go into one batch
- Use iter(loader) to walk the batches of one epoch
- The last batch can hold fewer rows than batch_size
- A DataLoader turns single rows from a Dataset into batches for the model
In this example, the DataLoader wraps the student dataset with a batch size of 2. The loop walks one epoch of batches. Each print shows the shapes the model would see. The first two batches hold 2 rows each. The last batch holds 1 row, because 5 rows do not split evenly by 2.
Try it Yourself
Edit the code, then press Run.
Running is simulated · Python Playground
Exercise
A DataLoader uses batch_size=2 on 5 rows. How many rows does the last batch hold?