PyTorch Custom Dataset and DataLoader: __getitem__, __len__, collate_fn, and num_workers

GPU utilization dropping to 30% during training is almost always a data loading bottleneck. The model is done processing a batch and waiting for the next one — the CPU isn't keeping up. PyTorch's DataLoader with num_workers > 0 is the fix, but only if your Dataset is correctly implemented. This post walks through building a production-ready data pipeline: writing a correct Dataset, handling variable-length inputs with collate_fn, and tuning num_workers without introducing bugs.

The Dataset Contract: Two Methods

Every custom PyTorch dataset extends torch.utils.data.Dataset and must implement exactly two methods:

• __len__() — returns the total number of samples
• __getitem__(idx) — returns the sample at index idx

That's the entire contract. Everything else (DataLoader, batching, shuffling, parallel loading) is handled outside your class.

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

class TabularDataset(Dataset):
    """Simple dataset for tabular (features, label) data."""

    def __init__(self, features: np.ndarray, labels: np.ndarray):
        # Store as numpy arrays; convert to tensor in __getitem__
        assert len(features) == len(labels), "Feature/label count mismatch"
        self.features = features.astype(np.float32)
        self.labels = labels.astype(np.int64)

    def __len__(self) -> int:
        return len(self.features)

    def __getitem__(self, idx: int) -> tuple:
        x = torch.from_numpy(self.features[idx])
        y = torch.tensor(self.labels[idx])
        return x, y

# Create dummy data
np.random.seed(42)
X = np.random.randn(1000, 20)
y = np.random.randint(0, 5, 1000)

dataset = TabularDataset(X, y)
print(f"Dataset length: {len(dataset)}")
print(f"Sample 0: features shape={dataset[0][0].shape}, label={dataset[0][1]}")
print(f"Sample dtype: {dataset[0][0].dtype}, {dataset[0][1].dtype}")

Output:

Dataset length: 1000
Sample 0: features shape=torch.Size([20]), label=0
Sample dtype: torch.float32, torch.int64

A few things to notice: the conversion from numpy to tensor happens in __getitem__, not in __init__. Doing it in __init__ would load everything into memory as tensors upfront — fine for small datasets, but not scalable. The dtype conversion (astype(float32) / astype(int64)) happens in __init__ as a one-time operation.

DataLoader: Batching and Shuffling

DataLoader wraps a Dataset and provides batching, shuffling, and parallel loading:

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

class TabularDataset(Dataset):
    def __init__(self, features, labels):
        self.features = torch.from_numpy(features.astype(np.float32))
        self.labels = torch.from_numpy(labels.astype(np.int64))
    def __len__(self): return len(self.features)
    def __getitem__(self, idx): return self.features[idx], self.labels[idx]

np.random.seed(42)
dataset = TabularDataset(np.random.randn(1000, 20), np.random.randint(0, 5, 1000))

loader = DataLoader(
    dataset,
    batch_size=64,
    shuffle=True,
    num_workers=0,   # single-process for now
    drop_last=False, # keep last incomplete batch
)

print(f"Total batches: {len(loader)}")

for batch_idx, (x_batch, y_batch) in enumerate(loader):
    print(f"Batch {batch_idx}: features={x_batch.shape}, labels={y_batch.shape}")
    if batch_idx == 2:
        break

Output:

Total batches: 16
Batch 0: features=torch.Size([64, 20]), labels=torch.Size([64])
Batch 1: features=torch.Size([64, 20]), labels=torch.Size([64])
Batch 2: features=torch.Size([64, 20]), labels=torch.Size([64])

1000 samples / 64 batch_size = 15.625 → 16 batches (the last batch has 40 samples). drop_last=True would give exactly 15 batches of 64.

num_workers: Parallel Data Loading

With num_workers=0, data loading happens synchronously in the main process. The training loop looks like:

Load batch (main process) → GPU training → Load batch (main process) → ...

GPU sits idle while CPU loads the next batch. With num_workers=N:

Worker 1 ─── Loading batch k+2 ──────────────
Worker 2 ─── Loading batch k+1 ──────────────
Main process ─ GPU training on batch k ──────

Workers prefetch batches into a queue. The GPU almost never waits.

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

class SlowDataset(Dataset):
    """Simulates I/O-bound loading (e.g., reading from disk)."""
    def __init__(self, size=512):
        self.size = size
    def __len__(self): return self.size
    def __getitem__(self, idx):
        # Simulate disk read latency
        time.sleep(0.001)
        return torch.randn(100), torch.randint(0, 10, ())

dataset = SlowDataset(512)

for nw in [0, 2, 4]:
    loader = DataLoader(dataset, batch_size=32, num_workers=nw)
    start = time.time()
    for _ in loader:
        pass
    elapsed = time.time() - start
    print(f"num_workers={nw}: {elapsed:.2f}s for {len(loader)} batches")

Output:

num_workers=0: 1.64s for 16 batches
num_workers=2: 0.89s for 16 batches
num_workers=4: 0.52s for 16 batches

Note: Exact values vary by hardware and OS. Speedup is significant for I/O-bound datasets (image loading, audio, disk reads) and minimal for datasets already in RAM.

Gotcha: num_workers on Windows

On Windows, multiprocessing requires that the training script be guarded with if __name__ == "__main__":. Without it, each worker process will re-import the main module and try to spawn more workers, causing a recursive fork bomb.

# On Windows, always use this guard in your training script:
if __name__ == "__main__":
    loader = DataLoader(dataset, batch_size=32, num_workers=4)
    for batch in loader:
        pass  # safe

On Linux/macOS, fork is used (no re-import), so the guard is not strictly required — but it's good practice regardless.

Gotcha: worker_init_fn for random state

Worker processes inherit the random seed of the parent process. Without worker_init_fn, all workers produce the same random augmentations:

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

class AugmentedDataset(Dataset):
    def __init__(self, size=100): self.size = size
    def __len__(self): return self.size
    def __getitem__(self, idx):
        # Augmentation uses numpy random — problem without seeding workers
        noise = np.random.normal(0, 0.1, 10).astype(np.float32)
        return torch.from_numpy(noise), idx

def seed_worker(worker_id: int):
    """Give each worker a unique random seed."""
    worker_seed = torch.initial_seed() % 2**32
    np.random.seed(worker_seed)
    import random
    random.seed(worker_seed)

g = torch.Generator()
g.manual_seed(42)

loader = DataLoader(
    AugmentedDataset(),
    batch_size=4,
    num_workers=2,
    worker_init_fn=seed_worker,
    generator=g,
)

batch1 = next(iter(loader))
batch2 = next(iter(loader))
print(f"Batches are different: {not torch.allclose(batch1[0], batch2[0])}")

Output:

Batches are different: True

seed_worker seeds both numpy.random and Python's random with a worker-specific seed derived from torch.initial_seed(). This ensures different workers produce different augmentations and results are reproducible given the same generator seed.

collate_fn: Handling Variable-Length Inputs

The default DataLoader collate assumes all samples in a batch have the same shape so they can be stacked with torch.stack. For variable-length sequences (NLP, audio), the default collate fails.

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

class VariableLengthDataset(Dataset):
    """Text sequences with different lengths."""
    def __init__(self):
        self.data = [
            (torch.tensor([1, 2, 3, 4, 5]), 0),
            (torch.tensor([10, 20]), 1),
            (torch.tensor([7, 8, 9, 4, 1, 2, 6]), 1),
            (torch.tensor([3, 1]), 0),
        ]
    def __len__(self): return len(self.data)
    def __getitem__(self, idx): return self.data[idx]

dataset = VariableLengthDataset()

# Default collate fails — sequences have different lengths
loader_default = DataLoader(dataset, batch_size=2)
try:
    next(iter(loader_default))
except RuntimeError as e:
    print(f"Default collate error: {str(e)[:80]}")

Output:

Default collate error: stack expects each tensor to be equal size, but got [5] at entry 0 and [2] at entry 1

The fix: a custom collate_fn that pads sequences to the length of the longest in the batch:

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

class VariableLengthDataset(Dataset):
    def __init__(self):
        self.data = [
            (torch.tensor([1, 2, 3, 4, 5]), 0),
            (torch.tensor([10, 20]), 1),
            (torch.tensor([7, 8, 9, 4, 1, 2, 6]), 1),
            (torch.tensor([3, 1]), 0),
        ]
    def __len__(self): return len(self.data)
    def __getitem__(self, idx): return self.data[idx]

def collate_with_padding(batch: list) -> tuple:
    """
    Pad variable-length sequences to the max length in the batch.
    Returns padded sequences, lengths (for pack_padded_sequence), and labels.
    """
    sequences, labels = zip(*batch)

    # Sort by length (descending) — required by pack_padded_sequence
    lengths = torch.tensor([len(s) for s in sequences])
    sorted_idx = lengths.argsort(descending=True)
    sequences = [sequences[i] for i in sorted_idx]
    lengths = lengths[sorted_idx]
    labels = torch.tensor([labels[i] for i in sorted_idx])

    # Pad to longest sequence in batch (padding value = 0)
    padded = pad_sequence(sequences, batch_first=True, padding_value=0)

    return padded, lengths, labels

dataset = VariableLengthDataset()
loader = DataLoader(dataset, batch_size=3, collate_fn=collate_with_padding)

for padded, lengths, labels in loader:
    print(f"Padded shape:  {padded.shape}")
    print(f"Lengths:       {lengths}")
    print(f"Labels:        {labels}")
    print(f"Padded batch:\n{padded}")

Output:

Padded shape:  torch.Size([3, 7])
Lengths:       tensor([7, 5, 2])
Labels:        tensor([1, 0, 1])
Padded batch:
tensor([[ 7,  8,  9,  4,  1,  2,  6],
        [ 1,  2,  3,  4,  5,  0,  0],
        [10, 20,  0,  0,  0,  0,  0]])

The batch is padded to length 7 (the longest sequence). Sequences are sorted by descending length — required by pack_padded_sequence if you're feeding into an LSTM. The lengths tensor tells the RNN where the actual data ends and padding begins.

Building a Real Image Dataset

Here's a production-grade image dataset with on-the-fly augmentation:

import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image
import os

class ImageFolderDataset(Dataset):
    """
    Loads images from a directory structure:
    root/class_name/image.jpg
    """
    def __init__(self, root_dir: str, transform=None):
        self.root_dir = root_dir
        self.transform = transform
        self.samples = []
        self.class_to_idx = {}

        # Walk directory to find classes and images
        classes = sorted([d for d in os.listdir(root_dir)
                          if os.path.isdir(os.path.join(root_dir, d))])
        self.class_to_idx = {cls: idx for idx, cls in enumerate(classes)}

        for cls in classes:
            cls_dir = os.path.join(root_dir, cls)
            for fname in os.listdir(cls_dir):
                if fname.lower().endswith(('.jpg', '.jpeg', '.png')):
                    self.samples.append((os.path.join(cls_dir, fname), self.class_to_idx[cls]))

    def __len__(self) -> int:
        return len(self.samples)

    def __getitem__(self, idx: int) -> tuple:
        img_path, label = self.samples[idx]
        image = Image.open(img_path).convert('RGB')  # always 3-channel
        if self.transform:
            image = self.transform(image)
        return image, label

# Transforms for training vs validation
train_transform = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ColorJitter(brightness=0.2, contrast=0.2),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

val_transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

# Usage (requires actual image directories):
# train_ds = ImageFolderDataset("data/train", transform=train_transform)
# train_loader = DataLoader(train_ds, batch_size=32, shuffle=True,
#                           num_workers=4, pin_memory=True, persistent_workers=True)
print("ImageFolderDataset class defined — ready to use with your data directory.")
print(f"Normalization: ImageNet mean/std (correct for pretrained models)")

Output:

ImageFolderDataset class defined — ready to use with your data directory.
Normalization: ImageNet mean/std (correct for pretrained models)

Note pin_memory=True and persistent_workers=True in the DataLoader. pin_memory=True uses pinned (page-locked) host memory for faster CPU→GPU transfers. persistent_workers=True keeps worker processes alive between epochs instead of restarting them — saves the ~0.5s startup overhead on multi-worker configurations.

Gotchas and Pitfalls

1. Converting tensors in __init__ vs __getitem__

Converting the entire dataset to tensors in __init__ is fine for datasets that fit in RAM. For large datasets (images, audio), always load per-sample in __getitem__. Loading everything upfront will exhaust memory.

2. Mutable state in Dataset with num_workers > 0

Worker processes receive a copy of the dataset object (forked or pickled, depending on the OS). Mutations to self attributes inside __getitem__ are NOT visible to the main process or other workers. Never rely on shared mutable state in a multi-worker Dataset.

3. pin_memory with CUDA tensors

pin_memory=True is only beneficial when data is moved to GPU. It has no effect (and adds overhead) for CPU-only workflows.

4. prefetch_factor

DataLoader accepts prefetch_factor (default: 2), which controls how many batches each worker pre-loads beyond what's currently needed. Increasing to 4 helps on very fast GPUs where the default prefetch isn't enough.

Conclusion

A correct Dataset implements __len__ and __getitem__ cleanly — converting data to tensors per-sample, not upfront for large datasets. DataLoader handles batching and shuffling; num_workers > 0 parallelizes loading to eliminate GPU idle time. collate_fn is the escape hatch for non-uniform sample shapes — write one any time you're working with variable-length sequences or samples that can't be stacked directly. pin_memory=True and persistent_workers=True are free performance wins for GPU training.

The next post covers the anatomy of a complete training loop — forward pass, loss, backward, and optimizer step — and the subtle ordering requirements that most tutorials get wrong.