Datasets and DataLoaders

What is a Dataset?

Deriving the Dataset Interface from First Principles

Start with the problem: We have N samples (images, texts, etc.). During training, we need to:

Access any sample by index (for batching)
Know the total count (for epoch iteration)

Naive approach: Load everything into a list.

data = [load_image(i) for i in range(N)]  # Uses N × image_size RAM!

Problem: For ImageNet (1.2M images, ~150GB), this crashes.

Solution: Lazy loading via indexing.

class Dataset:
    def __len__(self): return N
    def __getitem__(self, idx): return load_image(idx)  # Load on-demand

Now we load one image at a time, keeping RAM constant.

Example1: Custom Image Dataset from Scratch

Problem: Load images from a folder where filenames encode labels (e.g., cat_001.jpg).

import os
from PIL import Image
import torch
from torch.utils.data import Dataset
 
class CustomImageDataset(Dataset):
    def __init__(self, img_dir, transform=None):
        self.img_dir = img_dir
        self.transform = transform
        # List all image paths
        self.img_paths = [f for f in os.listdir(img_dir) if f.endswith('.jpg')]
        # WHY: We store paths, not images (lazy loading)
    def __len__(self):
        return len(self.img_paths)
        # WHY: Needed for batching and epoch calculation
    
    def __getitem__(self, idx):
        img_path = os.path.join(self.img_dir, self.img_paths[idx])
        image = Image.open(img_path)  # Load NOW (on-demand)
        # WHY: Only load when requested, not during __init__
        # Extract label from filename (e.g., "cat_001.jpg" → "cat")
        label = self.img_paths[idx].split('_')[0]
        if self.transform:
            image = self.transform(image)
            # WHY: Data augmentation (random crops, flips) happens per-access
        
        return image, label

Why this step (lazy loading)? If we loaded all images in __init__, we'd need gigabytes of RAM before training starts. Now we load ~1 image at a time.

What is a DataLoader?

Deriving DataLoader from Training Requirements

Training loop structure:

for epoch in range(num_epochs):
    for batch in data:  # Need batches, not individual samples
        loss = model(batch)
        loss.backward()

Requirements analysis:

Batching: SGD/Adam need batches (e.g., 32 samples) for gradient estimation.
- Math: Gradient estimate $\hat{g} = \frac{1}{B}\sum_{i=1}^B \nabla_\theta \mathcal{L}(x_i, y_i; \theta)$
- Why batches? Reduces gradient variance, enables GPU parallelism.
Shuffling: Prevents model from learning sample order (overfitting to sequence).
- Why each epoch? Model sees same data, but in different order → better generalization.
Paralelism: Data loading (disk I/O, decompression) is CPU-bound. GPU waits idle.
- Solution: Use multiple CPU processes to pre-load next batch while GPU trains on current.

Formula DataLoader Batching Mathematics

Given Dataset $\mathcal{D} = \{(x_i, y_i)\}_{i=1}^N$ and batch size $B$ :

\text{DataLoader}(\mathcal{D}, B) \to \left\{ \mathcal{B}_j \right\}_{j=1}^{\lceil N/B \rceil}

Where each batch:

\mathcal{B}_j = \{(x_i, y_i)\}_{i=(j-1)B+1}^{\min(jB, N)}

Collation (stacking):

\mathcal{B}_j \xrightarrow{\text{collate}} (X_j, Y_j) \in \mathbb{R}^{B \times d} \times \mathbb{R}^{B}

Derivation of number of batches:

Last batch may be incomplete: $|\mathcal{B}_{\text{last}}| = N \mod B$ (if $N \not\equiv 0 \mod B$ )
Option 1: drop_last=True → discard incomplete batch → $\lfloor N/B \rfloor$ batches
Option 2: drop_last=False → keep all → $\lceil N/B \rceil$ batches

Example 2: DataLoader with All Features

from torch.utils.data import DataLoader
 
# Assume we have the CustomImageDataset from before
dataset = CustomImageDataset(img_dir='./data/images')
 
dataloader = DataLoader(
    dataset,
    batch_size=32,          # Stack 32 samples per batch
    shuffle=True,           # Randomize order each epoch
    num_workers=4,          # 4 parallel CPU processes
    pin_memory=True,        # Speed up CPU → GPU transfer
    drop_last=True          # Drop incomplete last batch
)
 
# WHY each parameter:
# - batch_size=32: Trade-off between gradient accuracy and speed
#   (larger batch → less noise, but slower per-iteration)
# - shuffle=True: Breaks correlation between consecutive samples
#   (prevents overfitting to data order)
# - num_workers=4: While GPU trains batch t, CPUs load batch t+1
#   (overlap computation and I/O)
# - pin_memory=True: Allocates tensors in page-locked RAM
#   (faster DMA transfer to GPU, ~10-20% speedup)
# - drop_last=True: Ensures all batches have same size
#   (some models/layers require fixed batch size)

Using the DataLoader:

for epoch in range(10):
    for batch_idx, (images, labels) in enumerate(dataloader):
        # images: Tensor of shape (32, 3, 224, 224)
        # labels: Tensor of shape (32,)
        
        images = images.to(device)  # Move to GPU
        labels = labels.to(device)
        # Train step...
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

Why this step (enumerate)? We get both batch index (for logging) and data. The DataLoader internally calls dataset.__getitem__() 32 times, stacks results, and yields.

Deep Dive: Collation

Custom Collate Function

def collate_variable_length(batch):
    # batch: list of tuples [(seq1, label1), (seq2, label2), ...]
    # where seq_i has shape (length_i, feature_dim)
    # Step 1: Find max length
    max_len = max([seq.shape[0] for seq, _ in batch])
    # WHY: Tensors must be rectangular (all rows same length)
    
    # Step 2: Pad sequences to max_len
    padded_seqs = []
    labels = []
    for seq, label in batch:
        pad_amount = max_len - seq.shape[0]
        padded = torch.nn.functional.pad(seq, (0, 0, pad_amount))
        # WHY: Pad on length dimension (dim 0), not feature dimension (dim 1)
        padded_seqs.append(padded)
        labels.append(label)
    
    # Step 3: Stack into batch tensors
    batch_seqs = torch.stack(padded_seqs)    # (batch_size, max_len, feature_dim)
    batch_labels = torch.tensor(labels)      # (batch_size,)
    
    return batch_seqs, batch_labels
 
# Use with DataLoader
dataloader = DataLoader(dataset, batch_size=32, collate_fn=collate_variable_length)

Why padding? GPU operations (matrix multiplication) require fixed-size inputs. Padding wastes some computation but enables batching.

Common Patterns and Optimizations

Worked example Example 3: Data Augmentation in Dataset

Problem: Train image classifier with random crops and flips.

from torchvision import transforms
transform = transforms.Compose([
    transforms.RandomResizedCrop(224),      # Random crop and resize
    transforms.RandomHorizontalFlip(p=0.5), # 50% chance to flip
    transforms.ToTensor(),                   # Convert to tensor
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                        std=[0.229, 0.224, 0.225])  # ImageNet stats
])
dataset = CustomImageDataset(img_dir='./data', transform=transform)

Why apply transforms in Dataset? Each __getitem__() call gets a different augmentation (randomness per access). If we pre-applied, we'd get identical augmentations each epoch.

Worked example Example 4: Prefetching with Persistent Workers

Problem: High overhead from spawning new worker processes each epoch.

dataloader = DataLoader(
    dataset,
    batch_size=64,
    num_workers=8,
    persistent_workers=True,  # Keep workers alive across epochs
    prefetch_factor=2         # Each worker prefetches 2 batches
)

Why persistent_workers? Worker initialization (loading libraries, model state) takes seconds. With persistent_workers=True, this happens once per training session, not per epoch.

Why prefetch_factor=2? While GPU processes batch $t$ , workers load batches $t+1$ and $t+2$ . Hides I/O latency.

Diagram Explanation

The diagram shows the data pipeline flow:

Dataset: On-disk data (images, CSVs) with lazy __getitem__() access
Sampler: Generates indices (shuffled or sequential)
Workers: Parallel CPU processes fetch samples via indices
Collate: Stacks samples into batch tensors
GPU: Receives batched tensors for training

Key insight: The DataLoader is a producer-consumer system where CPU workers (producers) run ahead of GPU (consumer) to eliminate wait time.

Common mistake Common Mistake 1: Shuffling Inside Dataset

Wrong approach:

class MyDataset(Dataset):
    def __getitem__(self, idx):
        idx = random.randint(0, len(self) - 1)  # Random index
        return self.data[idx]

Why it feels right: "I want randomness, so I'll randomize the index."

Why it's wrong:

DataLoader shuffling already randomizes indices—you're doubling up
With multiple epochs, you might never see some samples (birthday paradox)
Breaks reproducibility (can't set seed reliably)

Fix: Let DataLoader handle shuffling via shuffle=True. Dataset should be deterministic.

Common mistake Common Mistake 2: Too Many Workers

Wrong approach:

dataloader = DataLoader(dataset, num_workers=32)  # "More is better!"

Why it feels right: More paralelism → faster loading.

Why it's wrong:

Spawning 32 processes has overhead (memory, context switching)
Disk I/O is the bottleneck, not CPU—32 workers fight for same disk bandwidth
Diminishing returns: 4-8 workers typically optimal

Fix: Benchmark with num_workers=0,2,4,8. Pick the value where training speed plateaus. Rule of thumb: num_workers =4 * num_GPUs.

Common mistake Common Mistake 3: Forgetting pin_memory for GPU Training

Wrong approach:

dataloader = DataLoader(dataset, batch_size=64)
for batch in dataloader:
    batch = batch.to('cuda')  # Slow transfer

Why it feels right: "I'll move data to GPU when needed."

Why it's wrong: Default tensors use pageable memory. OS may swap them to disk, causing unpredictable delays during .to('cuda').

Fix: Use pin_memory=True:

dataloader = DataLoader(dataset, batch_size=64, pin_memory=True)
for batch in dataloader:
    batch = batch.to('cuda', non_blocking=True)  # Fast DMA transfer

Why this works: Pinned (page-locked) memory can't be swapped, enabling faster, asynchronous GPU transfer.

Recall Explain to a 12-year-old

Imagine you're studying for a test with flashcards. You have 1000 cards in a big box.

Dataset is like organizing those cards. Each card has a number (index), and you can pull out any card by its number. But you DON'T pull out all1000 cards at once—that would cover your whole desk!

DataLoader is like your study buddy who:

Shuffles the cards so you don't memorize the order (that's cheating!)
Hands you 10 cards at a time (a batch) so you can study them together
Has 4 friends (workers) who pre-fetch the next batch while you study the current one—so you never wait

The magic: You study 1000 cards without ever having all of them on your desk at once, and your friends make sure the next batch is always ready!

Connections

3.3.01-Introduction-to-PyTorch: DataLoader is part of PyTorch's torch.utils.data module
3.3.05-Training-Loop-Basics: DataLoader feeds batches to the training loop
2.2.03-Mini-Batch-Gradient-Descent: Batch size parameter comes from DataLoader
3.2.04-Data-Augmentation: Transforms applied in Dataset's __getitem__()
4.1.02-Transfer-Learning: Pre-trained models expect specific normalization (applied in Dataset)
3.3.08-Distributed-Training: DataLoader with DistributedSampler for multi-GPU

#flashcards/ai-ml

What are the two methods a PyTorch Dataset must implement? :: __len__() (returns total samples) and __getitem__(idx) (returns sample at index idx)

Why do we use lazy loading in Datasets instead of loading all data in init?

Lazy loading loads data on-demand when getitem is called, keeping RAM usage constant. Loading all data at once would require N × sample_size RAM, which crashes for large datasets (e.g., ImageNet ~150GB).

What is the mathematical relationship between dataset size N, batch size B, and number of batches?

Number of batches = ⌈N/B⌉ if drop_last=False (keep incomplete batch), or ⌊N/B⌋ if drop_last=True (discard incomplete batch). The last batch size is N mod B when not dropping.

What does the DataLoader's collate function do?

It takes a list of individual samples from Dataset.getitem() and stacks them into batch tensors. Default uses torch.stack(), but custom collate functions handle variable-length data (e.g., padding sequences to max length).

Why use num_workers > 0 in DataLoader?

To parallelize data loading across CPU processes. While the GPU trains on batch t, CPU workers pre-load batch t+1, hiding I/O latency and preventing GPU from idling. Typical optimal value: 4-8 workers.

What does pin_memory=True do in DataLoader?

Allocates tensors in page-locked (pinned) RAM, which cannot be swapped to disk. This enables faster Direct Memory Access (DMA) for CPU-to-GPU transfer, providing ~10-20% speedup with non_blocking=True in .to('cuda').

Why should shuffling be done in DataLoader, not inside Dataset.getitem?

Dataset should be deterministic (same index → same sample) for reproducibility. DataLoader's shuffle=True randomizes indices before passing to Dataset, ensuring all samples are seen exactly once per epoch and maintaining seed-based reproducibility.

What is prefetch_factor in DataLoader? :: Number of batches each worker pre-loads ahead. With prefetch_factor=2 and num_workers=4, the system loads 8 batches ahead of current GPU batch, maximizing overlap between computation and I/O.

Why use drop_last=True in DataLoader?

Ensures all batches have identical size (batch_size). Some models/layers (e.g., BatchNorm with running statistics, certain distributed training setups) require or perform better with fixed batch sizes. Without it, the last batch may be smaller if N is not divisible by batch_size.

What is the time complexity of accessing a sample in a well-designed Dataset?

O(1) ideally, or O(log N) for indexed structures. Dataset should not iterate through all samples to find one. Store file paths or indices in a list for direct access by index.

Concept Map

Hinglish (regional understanding)

Intuition Hinglish mein samjho

Datasets aur DataLoaders ka sabse important kaam hai efficiently data feed karna model ko during training. Socho tumhare pas ImageNet jaisa dataset hai—millions of images. Sabko RAM me load karna impossible hai. Toh kya kare? Yahi pe Dataset aur DataLoader ka concept ata hai.

Datasetek blueprint hai jo bata hai ki tumhara data kahan hai aur kaise access karna hai. Jaise library me har book ka index hota hai, waise hi Dataset me har sample ka index hota hai. Jab tum kehte ho "mujhe sample number42 chahiye," tab Dataset disk se exactly woh sample load karta hai—lazy loading kehte hain isko. RAM me sab kuch ek sath nahi, bas jo chahiye woh load karo.

DataLoader ek smart librarian ki tarah kaam karta hai. Woh Dataset se samplesuthata hai, lekin efficiently. Pehle woh shuffles karta hai (taki model order yad na kare), phir batches banata hai (32 samples ek sath, kyunki GPU parallel processing me fast hai), aur multiple workers lagata hai (4-5

What is a Dataset?

Deriving the Dataset Interface from First Principles

Example1: Custom Image Dataset from Scratch

What is a DataLoader?

Deriving DataLoader from Training Requirements

Example 2: DataLoader with All Features

Deep Dive: Collation

Custom Collate Function

Common Patterns and Optimizations

Diagram Explanation

Connections

Concept Map

Hinglish (regional understanding)

Go deeper — visual, from zero

Test yourself — Deep Learning Frameworks

Connections