TopolearnReal data does not fit in one tensor; DataLoader is how PyTorch streams it efficiently into training.
torch.utils.data separates what a sample is from how samples are batched.
A map-style Dataset implements __len__ and __getitem__(i), returning one sample, often an (input, label) tuple. Load lazily in __getitem__ (read an image file, tokenize a line) so the dataset itself stays small. For tensors already in memory, TensorDataset(X, y) is enough, and random_split makes train/validation splits.
A DataLoader wraps a dataset and yields batches:
batch_sizeandshuffle=True(for training; keep validation unshuffled).num_workers=4loads batches in background processes so the GPU is not left waiting.pin_memory=Truespeeds up copies to CUDA;persistent_workers=Trueavoids restarting workers each epoch.drop_last=Truediscards a final short batch, useful when batch size must be constant.collate_fncontrols how samples are combined, for example padding variable-length sequences.
The default collate stacks tensors, numbers and NumPy arrays and preserves tuples and dicts, so for xb, yb in loader: just works.
With num_workers > 0 on Windows and macOS, workers are started with spawn, so the script must keep its training code under if __name__ == "__main__":. For data that cannot be indexed, such as a stream, subclass IterableDataset instead.
Example
import torch
from torch.utils.data import Dataset, DataLoader
class Squares(Dataset):
def __init__(self, n):
self.x = torch.arange(n, dtype=torch.float32)
def __len__(self):
return len(self.x)
def __getitem__(self, i):
return self.x[i], self.x[i] ** 2 # (input, target)
ds = Squares(10)
print(len(ds), ds[3]) # 10 (tensor(3.), tensor(9.))
loader = DataLoader(ds, batch_size=4, shuffle=True)
for xb, yb in loader:
print(xb.shape, yb.shape)
# torch.Size([4]) torch.Size([4])
# torch.Size([4]) torch.Size([4])
# torch.Size([2]) torch.Size([2])Remember
- Dataset = how to get sample i; DataLoader = batching, shuffling, workers
- Use
num_workersandpin_memoryto keep the GPU fed collate_fncustomises batching, e.g. padding
Watch out for
- Doing expensive loading in
__init__for huge datasets instead of lazily in__getitem__.