Data Pipelines: Dataset & DataLoader

1. The Two-Stage PyTorch Data Contract

In production deep learning, your training pipeline alternates between CPU-bound I/O (disk read, image decode, tokenization) and GPU-bound tensor math (matrix multiplication, backpropagation). PyTorch splits this responsibility:

2. Building a Custom Map-Style Dataset

Here is a production-ready custom dataset that loads numerical tabular features and target labels from an in-memory dictionary or CSV, applying on-the-fly normalization:

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

class CustomTabularDataset(Dataset):
    """
    Map-style Dataset subclass for tabular feature matrices.
    """
    def __init__(self, features: np.ndarray, labels: np.ndarray, normalize: bool = True):
        # Convert NumPy inputs to PyTorch tensors
        self.features = torch.tensor(features, dtype=torch.float32)
        self.labels = torch.tensor(labels, dtype=torch.long)
        
        # Optional feature normalization: (x - mean) / std
        if normalize:
            mean = self.features.mean(dim=0, keepdim=True)
            std = self.features.std(dim=0, keepdim=True) + 1e-7
            self.features = (self.features - mean) / std

    def __len__(self) -> int:
        """Returns the total number of samples."""
        return len(self.features)

    def __getitem__(self, idx: int):
        """Fetches and returns the (x, y) pair at sample index idx."""
        return self.features[idx], self.labels[idx]

# Mock data initialization (10,000 samples, 16 features, binary label)
mock_x = np.random.randn(10000, 16)
mock_y = np.random.randint(0, 2, size=(10000,))

dataset = CustomTabularDataset(mock_x, mock_y)
print(f"Total dataset length: {len(dataset)}")
sample_x, sample_y = dataset[42]
print(f"Sample 42 feature shape: {sample_x.shape}, label: {sample_y.item()}")

3. High-Throughput DataLoader Configuration

An improperly tuned DataLoader can starve your GPU, leaving expensive compute cores idle while waiting for disk reads. Below is the gold-standard production configuration:

# Production DataLoader Setup
train_loader = DataLoader(
    dataset=dataset,
    batch_size=64,
    shuffle=True,          # Reshuffle data at every epoch
    num_workers=4,         # Number of CPU worker processes fetching batches in background
    pin_memory=True,       # Allocates page-locked host memory for instant CUDA DMA transfer
    drop_last=True,        # Discard the final partial batch to preserve uniform batch size
    persistent_workers=True # Keeps background worker processes alive between epochs
)

# Iterating over the mini-batches
for batch_idx, (batch_features, batch_targets) in enumerate(train_loader):
    # Fast non-blocking transfer to GPU
    batch_features = batch_features.cuda(non_blocking=True)
    batch_targets = batch_targets.cuda(non_blocking=True)
    
    if batch_idx == 0:
        print(f"Batch shape: {batch_features.shape}, Targets: {batch_targets.shape}")
        break

4. Custom collate_fn for Variable-Length NLP Sequences

By default, DataLoader uses default_collate which expects all samples in a batch to have the exact same tensor dimensions. In Natural Language Processing (NLP) or Audio processing, sentences have varying word lengths. A custom collate_fn allows dynamic zero-padding:

import torch
from torch.nn.utils.rnn import pad_sequence

# Mock dataset returning variable-length integer token lists
class VariableTextDataset(Dataset):
    def __init__(self):
        self.samples = [
            torch.tensor([101, 2054, 2003, 1037, 102]),          # 5 tokens
            torch.tensor([101, 7592, 102]),                      # 3 tokens
            torch.tensor([101, 2023, 2003, 1037, 2742, 102]),    # 6 tokens
            torch.tensor([101, 2129, 102])                       # 3 tokens
        ]
        self.labels = torch.tensor([1, 0, 1, 0])

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

    def __getitem__(self, idx):
        return self.samples[idx], self.labels[idx]

def pad_collate_fn(batch):
    """
    Custom collator: receives a list of tuples [(x1, y1), (x2, y2), ...]
    and zero-pads features to match the longest sequence in this mini-batch.
    """
    sequences = [item[0] for item in batch]
    labels = torch.tensor([item[1] for item in batch], dtype=torch.long)
    
    # Pad sequences with token ID 0 (batch_first=True -> [B, Max_Seq_Len])
    padded_sequences = pad_sequence(sequences, batch_first=True, padding_value=0)
    
    # Compute attention mask (1 for real token, 0 for pad)
    attention_mask = (padded_sequences != 0).long()
    
    return padded_sequences, attention_mask, labels

nlp_dataset = VariableTextDataset()
nlp_loader = DataLoader(nlp_dataset, batch_size=2, shuffle=False, collate_fn=pad_collate_fn)

for batch_seq, batch_mask, batch_labels in nlp_loader:
    print(f"Padded Batch Shape: {batch_seq.shape}")
    print(f"Padded Sequences:\n{batch_seq}")
    print(f"Attention Mask:\n{batch_mask}\n")
    break

5. Common Traps & Multi-Worker Gotchas