Dataset & DataLoader Patterns

Separating Data Loading from Model Training

PyTorch decouples data loading from model logic using two classes:

1. torch.utils.data.Dataset: Encapsulates data fetching and preprocessing for a single sample.
2. torch.utils.data.DataLoader: Handles mini-batching, shuffling, parallel multi-threaded loading, and GPU memory pinning.

Raw Storage (Disk/S3) ──► [ Custom Dataset (__getitem__) ] ──► [ DataLoader (num_workers, batching) ] ──► GPU Mini-Batch

1. Implementing a Custom Dataset

A custom Dataset inherits from torch.utils.data.Dataset and must implement 3 Methods:

import torch
from torch.utils.data import Dataset


class CustomImageDataset(Dataset):

  def __init__(self, image_paths, labels, transform=None):
    self.image_paths = image_paths
    self.labels = labels
    self.transform = transform

  def __len__(self):
    # Return total sample count
    return len(self.image_paths)

  def __getitem__(self, idx):
    # Fetch sample and target label for index `idx`
    image_path = self.image_paths[idx]
    image = load_image(image_path)  # Custom loading function
    label = self.labels[idx]

    if self.transform:
      image = self.transform(image)

    return image, torch.tensor(label, dtype=torch.long)

2. Configuring High Performance DataLoader

from torch.utils.data import DataLoader

train_loader = DataLoader(
    dataset=train_dataset,
    batch_size=64,
    shuffle=True,  # Shuffle data every epoch
    num_workers=4,  # Use 4 parallel CPU subprocesses to fetch data
    pin_memory=True,  # Fast page-locked CPU memory to GPU transfer!
    drop_last=True,  # Drop incomplete final batch to maintain constant batch size
)

DataLoader Tuning Parameters

┌──────────────────────────┬──────────────────────────┬──────────────────────────┐
│ 1. NUM_WORKERS           │ 2. PIN_MEMORY            │ 3. COLLATE_FN            │
├──────────────────────────┼──────────────────────────┼──────────────────────────┤
│ Spawns N parallel CPU    │ Allocates page-locked CPU│ Custom function to pad   │
│ subprocesses to pre-fetch│ RAM memory, enabling fast│ variable length sequences│
│ data batches asynchronously.DMA transfer to CUDA GPUs! | into uniform batch tensors.│
└──────────────────────────┴──────────────────────────┴──────────────────────────┘

If num_workers = 0, data loading runs on the main Python thread. The GPU finishes computing a batch in 2ms, then sits idle for 50ms waiting for the CPU to load the next batch (GPU Starvation!).

Setting num_workers = 4 pre-fetches batches asynchronously so the GPU never waits.

Custom collate_fn for Variable Length Sequences

Default DataLoader expects all samples returned by __getitem__ to have identical tensor shapes.

For variable length text or audio sequences, supply a Custom collate_fn:

def pad_collate_fn(batch):
  # batch is a list of tuples [(text_tensor_1, label_1), (text_tensor_2, label_2)...]
  texts = [item[0] for item in batch]
  labels = [item[1] for item in batch]

  # Pad sequences dynamically to length of longest sequence in current batch
  padded_texts = torch.nn.utils.rnn.pad_sequence(
      texts, batch_first=True, padding_value=0
  )
  stacked_labels = torch.stack(labels)

  return padded_texts, stacked_labels

Say this out loud

PyTorch Dataset and DataLoader separate single sample fetching from batch iteration. Custom Datasets implement len and getitem methods. DataLoader handles multi-threaded parallel data loading, shuffling, dynamic padding via collate_fn, and pin_memory GPU transfers to prevent GPU starvation.

Followups to expect

1. What is IterableDataset in PyTorch? A dataset subclass used for streaming data sources (like Kafka streams or large S3 web tarballs) where random access index lookup via __getitem__ is impossible.
2. What is PyTorch Memory Leak in num_workers? Memory leaks caused by Python garbage collection issues when worker subprocesses copy large Python objects (like lists or dicts) in dataset initializers, resolved by using NumPy arrays or shared memory.