PyTorch Dataset and DataLoader
PyTorch Dataset and DataLoader separate "how to get one sample" from "how to feed batches to the model". A Dataset knows how many samples exist and how to return sample number i. A DataLoader wraps it to produce shuffled mini-batches, optionally using several worker processes to load data in parallel.
TensorDataset for In-Memory Data
import torch
from torch.utils.data import TensorDataset, DataLoader
X = torch.randn(100, 3)
y = torch.randint(0, 2, (100,))
ds = TensorDataset(X, y)
print(len(ds), ds[0][0].shape, ds[0][1].dtype)
loader = DataLoader(ds, batch_size=32, shuffle=True)
for xb, yb in loader:
print(xb.shape, yb.shape)100 torch.Size([3]) torch.int64 torch.Size([32, 3]) torch.Size([32]) torch.Size([32, 3]) torch.Size([32]) torch.Size([32, 3]) torch.Size([32]) torch.Size([4, 3]) torch.Size([4])
100 samples in batches of 32 gives three full batches and a final batch of 4. Pass drop_last=True if you want to discard the incomplete batch.
Writing a Custom Dataset
For files on disk, subclass Dataset and implement __len__ and __getitem__. Load lazily inside __getitem__ so you never hold everything in memory.
from torch.utils.data import Dataset
class SquaresDataset(Dataset):
def __init__(self, n):
self.n = n
def __len__(self):
return self.n
def __getitem__(self, idx):
x = torch.tensor([float(idx)])
return x, x ** 2
ds = SquaresDataset(10)
print(ds[3])
xb, yb = next(iter(DataLoader(ds, batch_size=4)))
print(xb.squeeze(1), yb.squeeze(1))(tensor([3.]), tensor([9.])) tensor([0., 1., 2., 3.]) tensor([0., 1., 4., 9.])
Splitting into Train and Validation
from torch.utils.data import random_split
full = SquaresDataset(1000)
train_ds, val_ds = random_split(full, [0.8, 0.2],
generator=torch.Generator().manual_seed(0))
print(len(train_ds), len(val_ds))800 200
Important DataLoader Options
- batch_size: samples per batch; larger is faster but uses more memory.
- shuffle=True for training, False for validation and testing.
- num_workers: background processes for loading; try 2-8 for image data. On Windows and macOS, guard your script with if __name__ == "__main__":.
- pin_memory=True speeds up copying batches to a CUDA GPU.
- collate_fn: custom function to combine samples, for example to pad variable-length sequences.
Transforms with torchvision
torchvision datasets accept a transform that runs on every sample, such as converting images to tensors and normalizing them:
from torchvision import datasets, transforms
tf = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,)),
])
mnist = datasets.MNIST(root="data", train=True, download=True, transform=tf)
img, label = mnist[0]
print(img.shape, img.dtype, label)torch.Size([1, 28, 28]) torch.float32 5
Frequently asked questions
What is the difference between Dataset and DataLoader in PyTorch?
A Dataset returns individual samples by index. A DataLoader iterates over a Dataset and groups samples into shuffled batches, optionally loading them in parallel.
How do I choose num_workers?
Start with 2-4 and increase while training speed improves. Use 0 when debugging, because errors in worker processes are harder to read.
Why does my DataLoader hang on Windows?
With num_workers above 0, Windows starts new processes that re-import your script. Put the training code inside if __name__ == "__main__": or set num_workers=0.