Dataset, DataLoader, and the Worker Plumbing

Two abstractions that keep your GPU fed

PyTorch separates "what is a sample" from "how do we batch them efficiently":

Dataset — defines __len__ and __getitem__(idx). The minimal contract: how many samples are there, and how do I get the i-th one.
DataLoader — wraps a Dataset and adds batching, shuffling, parallel loading via worker processes, and pin-memory for fast GPU transfer.

The DataLoader knobs that matter

batch_size — obvious.
shuffle=True — randomize order each epoch (training only; never on val/test).
num_workers — number of subprocesses for parallel loading. Set to N (where N = number of CPU cores) for I/O-bound datasets. 0 means main-process loading (slowest, but easiest to debug).
pin_memory=True — allocate batches in page-locked memory for faster CPU→GPU transfer. Combined with x.to(device, non_blocking=True), this gives real speedups.
prefetch_factor — how many batches each worker prefetches ahead. Default 2; raising helps when the GPU finishes a step before the next batch is ready.
persistent_workers=True — keep worker processes alive across epochs. Avoids the cost of re-spinning workers each epoch.
drop_last=True — drop the last incomplete batch. Useful when you want consistent batch sizes (BatchNorm stats, distributed training).

The macOS num_workers gotcha

On macOS, multiprocessing default is "spawn" (not "fork" as on Linux). This means each worker re-imports your module — heavy datasets can be slow to start, and any object that doesn't pickle cleanly will error. Workaround: use num_workers=0 for development, raise it for production runs.

Code

Custom Dataset — the minimal contract·python

import torch
from torch.utils.data import Dataset, DataLoader

class TensorDataset(Dataset):
    def __init__(self, X, y, transform=None):
        self.X = X
        self.y = y
        self.transform = transform

    def __len__(self):
        return len(self.X)

    def __getitem__(self, idx):
        x = self.X[idx]
        if self.transform is not None:
            x = self.transform(x)
        return x, self.y[idx]

X = torch.randn(1000, 10)
y = torch.randint(0, 5, (1000,))
ds = TensorDataset(X, y)
print(len(ds), ds[0][0].shape, ds[0][1])

Production-shaped DataLoader·python

import torch
from torch.utils.data import DataLoader

loader = DataLoader(
    dataset,
    batch_size=64,
    shuffle=True,
    num_workers=8,                # match CPU cores
    pin_memory=True,              # fast CPU→GPU
    prefetch_factor=4,            # batches each worker prefetches ahead
    persistent_workers=True,      # keep workers alive across epochs
    drop_last=True,               # consistent batch shape
)

device = "cuda"
for x, y in loader:
    x = x.to(device, non_blocking=True)
    y = y.to(device, non_blocking=True)
    # ...training step...

random_split — train / val from one Dataset·python

import torch
from torch.utils.data import random_split, DataLoader

dataset = TensorDataset(torch.randn(1000, 10), torch.randint(0, 5, (1000,)))
n_train = int(0.8 * len(dataset))
n_val = len(dataset) - n_train

train_ds, val_ds = random_split(
    dataset, [n_train, n_val],
    generator=torch.Generator().manual_seed(42),   # reproducible split
)

train_loader = DataLoader(train_ds, batch_size=32, shuffle=True)
val_loader = DataLoader(val_ds, batch_size=64, shuffle=False)

Custom collate — variable-length sequences·python

import torch
from torch.utils.data import DataLoader, Dataset
from torch.nn.utils.rnn import pad_sequence

class VariableSeqDataset(Dataset):
    def __init__(self):
        self.seqs = [torch.randint(0, 100, (torch.randint(5, 20, (1,)).item(),)) for _ in range(64)]
        self.labels = torch.randint(0, 2, (64,))

    def __len__(self): return len(self.seqs)
    def __getitem__(self, i): return self.seqs[i], self.labels[i]

def collate(batch):
    seqs, labels = zip(*batch)
    padded = pad_sequence(seqs, batch_first=True, padding_value=0)
    lengths = torch.tensor([len(s) for s in seqs])
    return padded, lengths, torch.tensor(labels)

loader = DataLoader(VariableSeqDataset(), batch_size=8, collate_fn=collate)
batch = next(iter(loader))
print(batch[0].shape, batch[1], batch[2].shape)
# torch.Size([8, max_len])  tensor([12, 18, ...])  torch.Size([8])

Dataset, DataLoader, and the Worker Plumbing

Two abstractions that keep your GPU fed

The DataLoader knobs that matter

The macOS num_workers gotcha

Code

External links

Exercise

Progress

Comments 0