tensorflow-data-pipelines

What this skill does

# TensorFlow Data Pipelines

Build efficient, scalable data pipelines using the tf.data API for optimal training performance. This skill covers dataset creation, transformations, batching, shuffling, prefetching, and advanced optimization techniques to maximize GPU/TPU utilization.

## Dataset Creation

### From Tensor Slices

```python
import tensorflow as tf
import numpy as np

# Create dataset from numpy arrays
x_train = np.random.rand(1000, 28, 28, 1)
y_train = np.random.randint(0, 10, 1000)

# Method 1: from_tensor_slices
dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))

# Apply transformations
dataset = dataset.shuffle(buffer_size=1024)
dataset = dataset.batch(32)
dataset = dataset.prefetch(tf.data.AUTOTUNE)

# Iterate through dataset
for batch_x, batch_y in dataset.take(2):
    print(f"Batch shape: {batch_x.shape}, Labels shape: {batch_y.shape}")
```

### From Generator Functions

```python
def data_generator():
    """Generator function for custom data loading."""
    for i in range(1000):
        # Simulate loading data from disk or API
        x = np.random.rand(28, 28, 1).astype(np.float32)
        y = np.random.randint(0, 10)
        yield x, y

# Create dataset from generator
dataset = tf.data.Dataset.from_generator(
    data_generator,
    output_signature=(
        tf.TensorSpec(shape=(28, 28, 1), dtype=tf.float32),
        tf.TensorSpec(shape=(), dtype=tf.int32)
    )
)

dataset = dataset.batch(32).prefetch(tf.data.AUTOTUNE)
```

### From Dataset Range

```python
# Create simple range dataset
dataset = tf.data.Dataset.range(1000)

# Use with custom mapping
dataset = dataset.map(lambda x: (tf.random.normal([28, 28, 1]), x % 10))
dataset = dataset.batch(32)
```

## Data Transformation

### Normalization Pipeline

```python
def normalize(image, label):
    """Normalize pixel values."""
    image = tf.cast(image, tf.float32) / 255.0
    return image, label

# Apply normalization
train_dataset = (
    tf.data.Dataset.from_tensor_slices((x_train, y_train))
    .map(normalize, num_parallel_calls=tf.data.AUTOTUNE)
    .batch(32)
    .prefetch(tf.data.AUTOTUNE)
)
```

### Data Augmentation Pipeline

```python
def augment(image, label):
    """Apply random augmentations."""
    image = tf.image.random_flip_left_right(image)
    image = tf.image.random_brightness(image, 0.2)
    image = tf.image.random_contrast(image, 0.8, 1.2)
    return image, label

def normalize(image, label):
    """Normalize pixel values."""
    image = tf.cast(image, tf.float32) / 255.0
    return image, label

# Build complete pipeline
train_dataset = (
    tf.data.Dataset.from_tensor_slices((x_train, y_train))
    .map(normalize, num_parallel_calls=tf.data.AUTOTUNE)
    .cache()  # Cache after normalization
    .shuffle(1000)
    .map(augment, num_parallel_calls=tf.data.AUTOTUNE)
    .batch(32)
    .prefetch(tf.data.AUTOTUNE)
)
```

### Multiple Transformations

```python
def resize_image(image, label):
    """Resize images to target size."""
    image = tf.image.resize(image, [224, 224])
    return image, label

def apply_random_rotation(image, label):
    """Apply random rotation augmentation."""
    angle = tf.random.uniform([], -0.2, 0.2)
    image = tfa.image.rotate(image, angle)
    return image, label

# Chain multiple transformations
dataset = (
    tf.data.Dataset.from_tensor_slices((images, labels))
    .map(resize_image, num_parallel_calls=tf.data.AUTOTUNE)
    .map(normalize, num_parallel_calls=tf.data.AUTOTUNE)
    .cache()
    .shuffle(10000)
    .map(augment, num_parallel_calls=tf.data.AUTOTUNE)
    .map(apply_random_rotation, num_parallel_calls=tf.data.AUTOTUNE)
    .batch(64)
    .prefetch(tf.data.AUTOTUNE)
)
```

## Batching and Shuffling

### Basic Batching Configuration

```python
# Batch size
BATCH_SIZE = 64

# Buffer size to shuffle the dataset
# (TF data is designed to work with possibly infinite sequences,
# so it doesn't attempt to shuffle the entire sequence in memory. Instead,
# it maintains a buffer in which it shuffles elements).
BUFFER_SIZE = 10000

dataset = dataset.shuffle(BUFFER_SIZE).batch(BATCH_SIZE, drop_remainder=True)
```

### Dynamic Batching

```python
# Variable batch sizes based on sequence length
def batch_by_sequence_length(dataset, batch_size, max_length):
    """Batch sequences by length for efficient padding."""
    def key_func(x, y):
        # Bucket by length
        return tf.cast(tf.size(x) / max_length * 10, tf.int64)

    def reduce_func(key, dataset):
        return dataset.batch(batch_size)

    return dataset.group_by_window(
        key_func=key_func,
        reduce_func=reduce_func,
        window_size=batch_size
    )
```

### Stratified Sampling

```python
def create_stratified_dataset(features, labels, batch_size):
    """Create dataset with balanced class sampling."""
    # Separate by class
    datasets = []
    for class_id in range(num_classes):
        mask = labels == class_id
        class_dataset = tf.data.Dataset.from_tensor_slices(
            (features[mask], labels[mask])
        )
        datasets.append(class_dataset)

    # Sample equally from each class
    balanced_dataset = tf.data.Dataset.sample_from_datasets(
        datasets,
        weights=[1.0/num_classes] * num_classes
    )

    return balanced_dataset.batch(batch_size).prefetch(tf.data.AUTOTUNE)
```

## Performance Optimization

### Caching Strategies

```python
# Cache in memory (for small datasets)
dataset = dataset.cache()

# Cache to disk (for larger datasets)
dataset = dataset.cache('/tmp/dataset_cache')

# Optimal caching placement
dataset = (
    tf.data.Dataset.from_tensor_slices((x_train, y_train))
    .map(expensive_preprocessing, num_parallel_calls=tf.data.AUTOTUNE)
    .cache()  # Cache after expensive operations
    .shuffle(buffer_size)
    .map(cheap_augmentation, num_parallel_calls=tf.data.AUTOTUNE)
    .batch(batch_size)
    .prefetch(tf.data.AUTOTUNE)
)
```

### Prefetching

```python
# Automatic prefetching
dataset = dataset.prefetch(tf.data.AUTOTUNE)

# Manual prefetch buffer size
dataset = dataset.prefetch(buffer_size=2)

# Complete optimized pipeline
optimized_dataset = (
    tf.data.Dataset.from_tensor_slices((x_train, y_train))
    .map(preprocess, num_parallel_calls=tf.data.AUTOTUNE)
    .cache()
    .shuffle(10000)
    .batch(64)
    .prefetch(tf.data.AUTOTUNE)
)
```

### Parallel Data Loading

```python
# Use num_parallel_calls for CPU-bound operations
dataset = dataset.map(
    preprocessing_function,
    num_parallel_calls=tf.data.AUTOTUNE
)

# Interleave for parallel file reading
def make_dataset_from_file(filename):
    return tf.data.TextLineDataset(filename)

filenames = tf.data.Dataset.list_files('/path/to/data/*.csv')
dataset = filenames.interleave(
    make_dataset_from_file,
    cycle_length=4,
    num_parallel_calls=tf.data.AUTOTUNE
)
```

### Memory Management

```python
# Use take() and skip() for train/val split without loading all data
total_size = 10000
train_size = int(0.8 * total_size)

full_dataset = tf.data.Dataset.from_tensor_slices((x, y))

train_dataset = (
    full_dataset
    .take(train_size)
    .shuffle(1000)
    .batch(32)
    .prefetch(tf.data.AUTOTUNE)
)

val_dataset = (
    full_dataset
    .skip(train_size)
    .batch(32)
    .prefetch(tf.data.AUTOTUNE)
)
```

## Advanced Patterns

### Iterating with For Loops

```python
# Basic iteration
for i in tf.data.Dataset.range(3):
    tf.print('iteration:', i)

# With dataset iterator
for i in iter(tf.data.Dataset.range(3)):
    tf.print('iteration:', i)
```

### Distributed Dataset

```python
# Distribute dataset across devices
for i in tf.distribute.OneDeviceStrategy('cpu').experimental_distribute_dataset(
    tf.data.Dataset.range(3)):
    tf.print('iteration:', i)

# Multi-GPU distribution
strategy = tf.distribute.MirroredStrategy()
distributed_dataset = strategy.experimental_distribute_dataset(dataset)
```

### Training Loop Integration

```python
# Execute training loop over dataset
for images, labels in tr