What are Dataset and DataLoader in PyTorch and how do they work together?

PyTorch's data pipeline follows a clean two-class design: Dataset defines how to access a single sample (index → data), and DataLoader wraps a Dataset to handle batching, shuffling, and parallel loading.

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

class TabularDataset(Dataset):
    def __init__(self, X: np.ndarray, y: np.ndarray):
        # Convert once at construction — not inside __getitem__!
        self.X = torch.tensor(X, dtype=torch.float32)
        self.y = torch.tensor(y, dtype=torch.long)

    def __len__(self) -> int:
        """Required — tells DataLoader how many samples exist."""
        return len(self.X)

    def __getitem__(self, idx: int):
        """Required — return a single (features, label) sample."""
        return self.X[idx], self.y[idx]

# Synthetic data
X = np.random.randn(1000, 20).astype(np.float32)
y = np.random.randint(0, 3, size=1000)

dataset = TabularDataset(X, y)
print(len(dataset))          # 1000
print(dataset[0])            # (tensor of 20 features, tensor scalar label)

loader = DataLoader(
    dataset,
    batch_size=32,
    shuffle=True,            # shuffle each epoch — essential for training
    num_workers=4,           # parallel data loading subprocesses
    pin_memory=True,         # faster CPU→GPU transfer
    drop_last=True,          # drop incomplete final batch
)

# Iterate over batches
for X_batch, y_batch in loader:
    print(X_batch.shape, y_batch.shape)  # (32, 20) (32,)
    break

# torchvision pre-built datasets
from torchvision import datasets, transforms
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,)),
])
mnist = datasets.MNIST(root="./data", train=True, download=True, transform=transform)
mnist_loader = DataLoader(mnist, batch_size=64, shuffle=True)