How does PyTorch's Dataset and DataLoader pipeline work, and what are the key performance considerations?

PyTorch's data loading follows a clean two-class design: Dataset encapsulates how to access a single sample (index → (X, y)), and DataLoader wraps a Dataset to handle batching, shuffling, and parallel data loading. Separating these responsibilities makes it easy to write dataset-specific logic once and reuse the same efficient loading infrastructure.

The most critical performance consideration is that the data loading pipeline must keep the GPU continuously fed — the GPU should never sit idle waiting for the next batch. Key knobs: num_workers launches subprocesses that prefetch batches in parallel with the GPU computation; pin_memory=True allocates batch tensors in pinned (non-pageable) CPU memory, enabling faster CPU→GPU transfers via DMA; prefetch_factor controls how many batches each worker prefetches ahead.

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 to tensors once at construction (not per __getitem__)
        self.X = torch.tensor(X, dtype=torch.float32)
        self.y = torch.tensor(y, dtype=torch.long)

    def __len__(self):
        return len(self.X)    # required — DataLoader uses this for indexing

    def __getitem__(self, idx):
        return self.X[idx], self.y[idx]  # single sample

dataset = TabularDataset(X_train, y_train)

loader = DataLoader(
    dataset,
    batch_size=256,
    shuffle=True,             # shuffle each epoch
    num_workers=4,            # parallel data loading
    pin_memory=True,          # faster CPU->GPU transfer
    drop_last=True,           # drop incomplete final batch
    persistent_workers=True,  # keep workers alive between epochs
)

# Training loop
for X_batch, y_batch in loader:
    X_batch = X_batch.cuda(non_blocking=True)  # async transfer
    y_batch = y_batch.cuda(non_blocking=True)
    # ... forward, backward, step