TensorFlow Guide: Working with Data - `tf.data` API

import tensorflow as tf
import numpy as np

# From a single tensor
dataset_single_tensor = tf.data.Dataset.from_tensor_slices([1, 2, 3, 4, 5])
for element in dataset_single_tensor:
    print(element.numpy())

print("\n---")

# From multiple tensors (features and labels)
features = tf.constant([[10, 20], [30, 40], [50, 60]])
labels = tf.constant([0, 1, 0])
dataset_multi_tensor = tf.data.Dataset.from_tensor_slices((features, labels)) # Tuple of tensors

for x, y in dataset_multi_tensor:
    print(f"Features: {x.numpy()}, Label: {y.numpy()}")

print("\n---")

# From NumPy arrays
numpy_features = np.array([[1.1, 1.2], [2.1, 2.2]])
numpy_labels = np.array([0, 1])
dataset_numpy = tf.data.Dataset.from_tensor_slices((numpy_features, numpy_labels))

for x, y in dataset_numpy:
    print(f"NumPy Features: {x.numpy()}, NumPy Label: {y.numpy()}")

import tensorflow as tf

# A simple Python generator function
def data_generator():
    for i in range(1, 6):
        yield i, i * 10 # Yield (feature, label) pairs

# Create a dataset from the generator
# output_types specifies the data types of the yielded elements
# output_shapes specifies the shape of the yielded elements (None for dynamic shapes)
dataset_from_generator = tf.data.Dataset.from_generator(
    data_generator,
    output_types=(tf.int32, tf.int32),
    output_shapes=(None, None) # Or (tf.TensorShape([]), tf.TensorShape([])) for scalars
)

print("\n--- Dataset from Generator ---")
for x, y in dataset_from_generator:
    print(f"Generated X: {x.numpy()}, Generated Y: {y.numpy()}")

import tensorflow as tf

dataset = tf.data.Dataset.from_tensor_slices([1, 2, 3, 4, 5])

def square_and_add_one(x):
    return tf.square(x) + 1

# Apply the function to each element
mapped_dataset = dataset.map(square_and_add_one)

print("\n--- Mapped Dataset (square_and_add_one) ---")
for element in mapped_dataset:
    print(element.numpy())

# With multiple inputs/outputs
features = tf.constant([1.0, 2.0, 3.0])
labels = tf.constant([0, 1, 0])
dataset_pair = tf.data.Dataset.from_tensor_slices((features, labels))

def preprocess_pair(x, y):
    x_scaled = x * 10 # Scale feature
    y_one_hot = tf.one_hot(y, depth=2) # One-hot encode label (assuming 2 classes)
    return x_scaled, y_one_hot

processed_dataset = dataset_pair.map(preprocess_pair)

print("\n--- Mapped Dataset (processed pairs) ---")
for x, y in processed_dataset:
    print(f"Processed X: {x.numpy()}, Processed Y: {y.numpy()}")

# For images, map is used for resizing, augmentation, etc.
# For optimal performance, use num_parallel_calls=tf.data.AUTOTUNE with map()
# This parallelizes data processing across multiple CPU cores.
# processed_dataset = dataset_pair.map(preprocess_pair, num_parallel_calls=tf.data.AUTOTUNE)

import tensorflow as tf

dataset = tf.data.Dataset.from_tensor_slices([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])

def is_even(x):
    return x % 2 == 0

filtered_dataset = dataset.filter(is_even)

print("\n--- Filtered Dataset (even numbers) ---")
for element in filtered_dataset:
    print(element.numpy())

import tensorflow as tf

dataset = tf.data.Dataset.from_tensor_slices([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])

# buffer_size should be larger than your dataset for perfect shuffling,
# or large enough to provide a good level of randomness.
shuffled_dataset = dataset.shuffle(buffer_size=10)

print("\n--- Shuffled Dataset ---")
for element in shuffled_dataset:
    print(element.numpy())

import tensorflow as tf

dataset = tf.data.Dataset.from_tensor_slices([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])

batched_dataset = dataset.batch(batch_size=3)

print("\n--- Batched Dataset ---")
for batch in batched_dataset:
    print(batch.numpy())

# Handling remaining elements: drop_remainder=True will discard batches smaller than batch_size
batched_dataset_dr = dataset.batch(batch_size=3, drop_remainder=True)
print("\n--- Batched Dataset (drop_remainder=True) ---")
for batch in batched_dataset_dr:
    print(batch.numpy())

import tensorflow as tf

dataset = tf.data.Dataset.from_tensor_slices([1, 2])

repeated_dataset = dataset.repeat(count=3) # Repeats 3 times

print("\n--- Repeated Dataset ---")
for element in repeated_dataset:
    print(element.numpy())

import tensorflow as tf
import time

# Simulate slow data loading and processing
def slow_data_generator():
    for i in range(5):
        time.sleep(0.1) # Simulate IO delay
        yield i

def slow_preprocessing(x):
    time.sleep(0.05) # Simulate CPU processing delay
    return x * 2

# Without prefetching
dataset_no_prefetch = tf.data.Dataset.from_generator(
    slow_data_generator, output_types=tf.int32
).map(slow_preprocessing).batch(2)

start_time = time.time()
print("\n--- Without Prefetching ---")
for batch in dataset_no_prefetch:
    print(f"Processed batch: {batch.numpy()}")
print(f"Time without prefetch: {time.time() - start_time:.4f} seconds\n")


# With prefetching
dataset_with_prefetch = tf.data.Dataset.from_generator(
    slow_data_generator, output_types=tf.int32
).map(slow_preprocessing, num_parallel_calls=tf.data.AUTOTUNE).batch(2).prefetch(tf.data.AUTOTUNE)

start_time = time.time()
print("\n--- With Prefetching ---")
for batch in dataset_with_prefetch:
    print(f"Processed batch: {batch.numpy()}")
print(f"Time with prefetch: {time.time() - start_time:.4f} seconds")

import tensorflow as tf
import numpy as np

# Simulate some large data (e.g., image paths and labels)
NUM_SAMPLES = 10000
image_paths = [f'/data/img_{i}.png' for i in range(NUM_SAMPLES)]
labels = np.random.randint(0, 10, size=NUM_SAMPLES) # 10 classes

# 1. Create a dataset of (image_path, label) pairs
dataset = tf.data.Dataset.from_tensor_slices((image_paths, labels))

# 2. Shuffle early for better randomness
BUFFER_SIZE = 5000 # Larger buffer for better shuffling
dataset = dataset.shuffle(buffer_size=BUFFER_SIZE)

# Simulate image loading and preprocessing (this would be where you load from disk)
def load_and_preprocess_image(image_path, label):
    # In a real scenario, you'd use tf.io.read_file, tf.image.decode_jpeg, etc.
    # For this example, let's just create a dummy tensor
    # Simulate a small delay for IO and computation
    # tf.py_function is used to wrap arbitrary Python functions for tf.data
    # This is useful for functions that cannot be expressed purely in TensorFlow ops.
    def _load_and_preprocess_py(path, lbl):
        # time.sleep(0.01) # Simulate I/O or heavy CPU work
        dummy_image = np.random.rand(64, 64, 3).astype(np.float32) # Dummy 64x64 RGB image
        return dummy_image, lbl

    image, new_label = tf.py_function(
        _load_and_preprocess_py,
        inp=[image_path, label],
        Tout=[tf.float32, tf.int64] # Specify output types
    )
    # Important: set the shape after py_function as it loses static shape info
    image.set_shape([64, 64, 3])
    new_label.set_shape([]) # Scalar label
    return image, new_label

# 3. Map preprocessing function, parallelize with AUTOTUNE
dataset = dataset.map(load_and_preprocess_image, num_parallel_calls=tf.data.AUTOTUNE)

# 4. (Optional) Cache if data fits in memory and mapping is expensive
# dataset = dataset.cache()

# 5. Batch the data
BATCH_SIZE = 32
dataset = dataset.batch(BATCH_SIZE)

# 6. Prefetch to overlap data production and consumption
dataset = dataset.prefetch(tf.data.AUTOTUNE)

# Now you can iterate over this `dataset` in your training loop
print("\n--- Optimized Pipeline Example ---")
for i, (images, labels) in enumerate(dataset.take(3)): # Take 3 batches for demonstration
    print(f"Batch {i+1}: Images shape={images.shape}, Labels shape={labels.shape}")

# Your solution for Exercise 4.1 here
import tensorflow as tf
import numpy as np
import time

# 1. Simulate Data
NUM_SAMPLES = 1000
image_paths_dummy = [f'dummy_path_{i}.jpg' for i in range(NUM_SAMPLES)]
labels_dummy = np.random.randint(0, 10, size=NUM_SAMPLES) # 10 classes

# Create a tf.data.Dataset from dummy paths and labels
raw_dataset = tf.data.Dataset.from_tensor_slices((image_paths_dummy, labels_dummy))

# 2. Preprocessing Function
# This function would typically load an actual image from disk
# and perform operations like resizing, augmentation, etc.
def process_image(image_path, label):
    # Simulate loading an image by creating a random tensor
    # In a real scenario:
    # img_raw = tf.io.read_file(image_path)
    # img_tensor = tf.image.decode_jpeg(img_raw, channels=3)
    img_tensor = tf.random.uniform(shape=(128, 128, 3), minval=0.0, maxval=255.0, dtype=tf.float32)

    # Normalize pixel values to 0-1
    img_normalized = img_tensor / 255.0

    # Apply random horizontal flipping
    img_augmented = tf.image.random_flip_left_right(img_normalized)

    return img_augmented, label

# Due to `tf.image.random_flip_left_right` being a TF op, `tf.py_function` might not be strictly necessary
# if all operations are pure TensorFlow. However, if you had more complex Python logic, you would use it.
# Let's write it purely in TensorFlow for optimal performance.

# 3. Build the Pipeline
BUFFER_SIZE = 500
BATCH_SIZE = 64

# Shuffle early
dataset_pipeline = raw_dataset.shuffle(buffer_size=BUFFER_SIZE)

# Map preprocessing function, parallelize
dataset_pipeline = dataset_pipeline.map(process_image, num_parallel_calls=tf.data.AUTOTUNE)

# Batch the data
dataset_pipeline = dataset_pipeline.batch(BATCH_SIZE)

# Prefetch to overlap data production and consumption
dataset_pipeline = dataset_pipeline.prefetch(tf.data.AUTOTUNE)

# 4. Verify the pipeline
print("\n--- Verifying Custom Image Data Pipeline ---")
for i, (images, labels) in enumerate(dataset_pipeline.take(3)): # Take 3 batches for verification
    print(f"Batch {i+1}:")
    print(f"  Images shape: {images.shape}")
    print(f"  Images dtype: {images.dtype}")
    print(f"  Labels shape: {labels.shape}")
    print(f"  Labels dtype: {labels.dtype}")
    # Assertions to ensure correct shapes and types
    assert images.shape == (BATCH_SIZE, 128, 128, 3), "Image batch shape is incorrect!"
    assert labels.shape == (BATCH_SIZE,), "Label batch shape is incorrect!"
    assert images.dtype == tf.float32, "Image dtype is incorrect!"
    assert labels.dtype == tf.int64 or labels.dtype == tf.int32, "Label dtype is incorrect!"

print("\nPipeline verified successfully for a few batches!")