data pipeline ml
**ML Data Pipeline** is the **system that efficiently loads, preprocesses, and batches training data** — a bottleneck that can reduce GPU utilization from 100% to < 30% if poorly implemented, making data loading optimization as important as model architecture.
**The I/O Bottleneck Problem**
- GPU throughput: Processes a batch in 50ms.
- Naive data loading: Read from disk + decode + augment = 200ms per batch.
- Result: GPU idle 75% of the time — $3,000/month GPU cluster at 25% utilization.
- Solution: Overlap data preparation with GPU compute using prefetching and parallel loading.
**PyTorch DataLoader**
```python
dataloader = DataLoader(
dataset,
batch_size=256,
num_workers=8, # Parallel CPU workers
prefetch_factor=2, # Batches to prefetch per worker
pin_memory=True, # Pinned memory for fast GPU transfer
persistent_workers=True # Avoid worker restart overhead
)
```
- `num_workers`: Spawn N CPU processes for parallel loading. Rule of thumb: 4× number of GPUs.
- `prefetch_factor`: Each worker prefetches factor× batches ahead.
- `pin_memory=True`: Required for async GPU transfer.
**TensorFlow `tf.data` Pipeline**
```python
dataset = tf.data.Dataset.from_tensor_slices(filenames)
dataset = dataset.interleave(tf.data.TFRecordDataset, num_parallel_calls=8)
dataset = dataset.map(preprocess, num_parallel_calls=tf.data.AUTOTUNE)
dataset = dataset.batch(256)
dataset = dataset.prefetch(tf.data.AUTOTUNE) # Overlap GPU compute with CPU prep
```
**Storage Optimization**
- **TFRecord / WebDataset**: Sequential binary format → faster disk reads than random file access.
- **LMDB**: Memory-mapped key-value store — near-RAM speeds for small datasets.
- **Petastorm**: Distributed dataset format for Spark + PyTorch/TF.
**Online Augmentation**
- Apply augmentations (crop, flip, color jitter) on CPU workers during loading — free compute.
- GPU augmentation (NVIDIA DALI): Move decode and augment to GPU — further reduces CPU bottleneck.
Efficient data pipeline design is **a critical ML engineering skill** — well-tuned data loading routinely improves training throughput 2-5x with no changes to model architecture, directly reducing the cost and time of every training run.