Designing Data Pipelines with TensorFlow

Table of Contents

• Module 1: Evaluating TensorFlow Capabilities
  • What’s New in TensorFlow 2.0
  • Migrating from TensorFlow 1.0 to TensorFlow 2.0
  • Data Pipeline and Model Training Overview
• Module 2: Loading Data in TensorFlow
  • The tf.data.Dataset Object
  • Loading Data from a CSV
  • Loading Data from NumPy and pandas
  • Serializing Data with tf.Example
  • Storing Data Efficiently with TFRecord
  • Loading Image Data
• Module 3: Prepping Data
  • Feature Engineering with pandas
  • Combining Datasets with zip and map
  • Loading Image Data with Generators
  • Image Data Augmentation
• Module 4: Optimizing Performance of Pipelines
  • Preparing Data for Model Training
  • Building Models with the Keras Sequential API
  • Batching and Prefetching
  • Parallelizing Data Extraction with interleave
• Summary

This course covers one of the major improvements introduced in the TensorFlow 2.0 release: the tf.data module. This module brings a number of performance improvements that make it much easier to build efficient models in TensorFlow. Topics covered include migrating from TensorFlow 1.0 to TensorFlow 2.0, loading data into tf.data.Dataset objects, prepping data and feature engineering, and — most importantly — optimizing data pipeline performance. By the end of this material, you should know how to use the tf.data module in TensorFlow and how to take advantage of many of its performance improvements.

Prerequisites: familiarity with Python and a basic understanding of TensorFlow. Knowledge of the Keras API is helpful but not required.

Module 1: Evaluating TensorFlow Capabilities

This module covers the major improvements introduced in TensorFlow 2.0 and specifically what changed with respect to tf.data, the object at the center of this course on designing efficient data pipelines. It also covers how to migrate code from TensorFlow 1.0 to TensorFlow 2.0, as well as a general overview of how a dataset fits together with Keras. By the end of this module, the stage is set for building efficient data pipelines. If you are brand new to TensorFlow and have never used TensorFlow 1.0, you can skip ahead to Module 2, since that is where the pipeline-building content truly begins.

What’s New in TensorFlow 2.0

TensorFlow is a relatively young framework compared to some of its predecessors, but it took over the field thanks to an approachable API and interface. Google, the primary backer of TensorFlow, made several significant improvements in the 2.0 release. These contributions fall into a few major buckets:

flowchart TB
    TF["TensorFlow 2.0 Improvements"] --> A["Easier Model Building"]
    TF --> B["Robust Deployment (TFX)"]
    TF --> C["Consolidated Utilities"]
    TF --> D["Improved Data Pipelines (tf.data)"]

    A --> A1["Keras API integrated and simplified"]
    A --> A2["Eager execution by default"]

    B --> B1["TFX framework for production deployment"]

    C --> C1["Previously custom internal libraries now built in"]

    D --> D1["Faster data loading from disk"]
    D --> D2["Easier splitting of data"]
    D --> D3["Parallelization of data loading"]

• Easier model building — This encompasses the combination and dramatic simplification of the Keras API, making it easy to layer on model-building elements, along with eager execution, which greatly improves the ability to combine model objects and generate quality outputs.
• Robust deployment into production — The TFX (TensorFlow Extended) framework is a significant addition and is worth exploring separately if production deployment of TensorFlow objects is a goal.
• Consolidated utilities — Previously, many teams built their own internal libraries around TensorFlow. Google has assimilated a lot of those utilities directly under the TensorFlow umbrella.
• Improved data pipelines — This is the focus of this course. The tf.data module does a lot to help build quality models efficiently: it gets data off disk faster, makes splitting data easier, and allows parallelizing data loading.

None of this existed in TensorFlow 1.0. There were workarounds, but they required a solid understanding of the internals to use effectively. With tf.data, a data scientist can leverage these improvements without needing to understand TensorFlow’s internals in depth. Benefits include:

• Simpler deployment — Data objects can be saved in an easier fashion, and it becomes straightforward to ensure that inference-time inputs match what was seen during training.
• A simplified API — The API is approachable even for practitioners who are not hardcore engineers.
• Performance optimizations — With the advent of GPUs and TPUs increasing the speed of model training, the data input pipeline has become the new bottleneck. tf.data includes utilities specifically designed to get data into place for training much faster — especially important when data lives in a different location from where training happens.

Migrating from TensorFlow 1.0 to TensorFlow 2.0

A common question when evaluating TensorFlow 2.0 is whether code written for TensorFlow 1.0 will continue to work. In most cases, if the original code performs detailed transformations, the honest answer is probably not without modification. However, simplified code — for example, code that primarily manipulates data with pandas and does not lean heavily on TensorFlow-specific transformations — may work without changes.

The catch: even when old code runs, it will not be able to take advantage of the real benefits of TensorFlow 2.0, especially tf.data.

In many cases, TensorFlow 2.0 is backward compatible, meaning it is possible to run certain code without modification. One way to achieve this is by globally forcing TensorFlow 2.0 to behave like TensorFlow 1.0:

import tensorflow as tf

tf.compat.v1.disable_v2_behavior()

The major downside of doing this is that it also disables access to the tf.data portion of TensorFlow 2.0 — arguably one of the biggest benefits of upgrading. This approach might make sense if existing data pipelines are already established and only the Keras API needs to be reused, but it forfeits the tf.data improvements, which matter a great deal when data loading is the bottleneck.

TensorFlow’s official documentation contains a detailed migration guide covering specifics and gotchas, and TensorFlow also ships a migration script that can be downloaded and run against existing code to automatically update it for TensorFlow 2.0 compatibility. It does not work in every case, but it does handle many common patterns.

Data Pipeline and Model Training Overview

Before diving into pipeline construction, it helps to review the general Keras model-training process end to end, since later modules build directly on top of this flow.

flowchart LR
    A["Load libraries\n(tf, numpy, keras)"] --> B["Load dataset\n(fashion_mnist.load_data)"]
    B --> C["Inspect shape\n(60000, 28, 28)"]
    C --> D["Normalize pixel values\n0-255 -> 0-1"]
    D --> E["Build model\n(Sequential API)"]
    E --> F["Compile model\n(optimizer, loss, metrics)"]
    F --> G["Fit model\n(train_images, train_labels, epochs)"]
    G --> H["Predict\n(model.predict)"]

The following example uses the Fashion-MNIST dataset, which is similar to the classic MNIST character-recognition dataset but uses articles of clothing instead of digits:

import tensorflow as tf
import numpy as np
from tensorflow import keras

# Load the Fashion-MNIST dataset
fashion_mnist = keras.datasets.fashion_mnist
(train_images, train_labels), (test_images, test_labels) = fashion_mnist.load_data()

# Inspect the shape of the training data
print(train_images.shape)
# (60000, 28, 28) -> 60,000 observations, each a 28x28 pixel image

# Normalize pixel values from the 0-255 range down to 0-1
train_images = train_images / 255.0
test_images = test_images / 255.0

# Build the model using the Keras Sequential API
model = keras.Sequential([
    keras.layers.Flatten(input_shape=(28, 28)),
    keras.layers.Dense(128, activation="relu"),
    keras.layers.Dense(10)
])

# Compile the model
model.compile(
    optimizer="adam",
    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=["accuracy"]
)

# Fit the model
model.fit(train_images, train_labels, epochs=10)
# ... final epoch reports accuracy: 0.9098 (over 90% with only 10 epochs on CPU)

# Generate predictions
predictions = model.predict(test_images)
print(np.argmax(predictions[0]))  # index of the most likely class for the first test image
print(test_labels[0])             # compare against the actual label

Key points from this walkthrough:

• The training images have a shape of (60000, 28, 28) — 60,000 observations, each with two spatial dimensions (28x28 pixels), giving an effectively three-dimensional NumPy array.
• Pixel values range from 0-255 and must be scaled down to the 0-1 range before being fed into the model, since neural networks generally expect normalized inputs.
• The Sequential API layers are added in order: Flatten (required for image data, not for tabular data) with the input_shape matching the image dimensions, a Dense hidden layer (128 nodes is common for smaller image datasets, with a relu activation), and a final Dense output layer with one node per class (10 classes for Fashion-MNIST).
• Compiling requires specifying an optimizer (adam), a loss function (SparseCategoricalCrossentropy with from_logits=True since the output layer produces raw logits), and a metric to monitor (accuracy).
• model.fit mirrors the fit APIs found in libraries like scikit-learn: pass in X (train_images), y (train_labels), and the number of epochs.
• model.predict returns logit arrays; taking the argmax of each row gives the predicted class index, which can then be compared against the true test label.

This example already assumes a dataset has been built. The rest of the course covers exactly how to construct, prep, and optimize that dataset object.

Module 2: Loading Data in TensorFlow

This module introduces the tf.data.Dataset type and shows how to load data in multiple different formats. By the end of it, data can move in and out of TensorFlow using the tf.data.Dataset type in order to unlock the performance optimizations covered later.

The tf.data.Dataset Object

tf.data.Dataset is the fundamental component for optimizing pipelines inside TensorFlow. It is possible to train Keras models without wrapping data in a dataset object, but doing so limits what performance optimizations are available. When working with larger data or streaming data — where loading data into the training pipeline is the bottleneck — data should be stored inside one of these dataset objects.

There are two general ways to construct a tf.data.Dataset:

flowchart TB
    Start["Need a tf.data.Dataset"] --> Q{"Where does the data\ncome from?"}
    Q -->|"Already a prebuilt dataset\n(in memory or on disk)"| A["From a data source\ne.g. dataset shipped in docs,\nprovided by another team"]
    Q -->|"Existing data in another\nformat/object"| B["From a data transformation\nCSV, pandas, NumPy, etc.\nconverted into a Dataset"]
    A --> Result["tf.data.Dataset"]
    B --> Result

1. From a data source — A tf.data.Dataset constructed and stored in memory or a file that can simply be loaded. This is the simplest option, such as loading an established dataset from documentation or from a dataset provided by another team.
2. From a data transformation — Constructing a tf.data.Dataset from other data objects: CSV files, pandas DataFrames, NumPy arrays, and other data types can all be converted into a dataset. This is the more common path in real projects.

This module covers loading data from a CSV, from NumPy and pandas, and then the big performance-oriented formats: tf.Example and TFRecord. It finishes with loading image data.

Loading Data from a CSV

CSV is one of the most common data storage mediums for transferring data. Loading a CSV into a tf.data.Dataset uses the experimental CSV utility built into tf.data:

import numpy as np
import pandas as pd
import tensorflow as tf

TRAIN_DATA_URL = "https://storage.googleapis.com/tf-datasets/titanic/train.csv"
TEST_DATA_URL = "https://storage.googleapis.com/tf-datasets/titanic/eval.csv"

# Download the CSV files and get their local paths
train_file_path = tf.keras.utils.get_file("train.csv", TRAIN_DATA_URL)
test_file_path = tf.keras.utils.get_file("eval.csv", TEST_DATA_URL)
# e.g. ~/.keras/datasets/train.csv

LABEL_COLUMN = "survived"
LABELS = [0, 1]

def get_dataset(file_path, **kwargs):
    dataset = tf.data.experimental.make_csv_dataset(
        file_path,
        batch_size=5,          # small batch size purely to make the data easy to inspect
        label_name=LABEL_COLUMN,
        na_value="?",
        num_epochs=1,
        ignore_errors=True,
        **kwargs
    )
    return dataset

raw_train_data = get_dataset(train_file_path)
raw_test_data = get_dataset(test_file_path)

# raw_test_data is a PrefetchDataset -- data is fetched as needed to keep training moving

def show_batch(dataset):
    for batch, label in dataset.take(1):
        for key, value in batch.items():
            print(f"{key:20s}: {value.numpy()}")

show_batch(raw_train_data)
# Columns picked up automatically from the CSV header:
# sex, age, n_siblings_spouses, parch, fare, class, deck, embark_town, alone

If a CSV has no header row, column names can be supplied explicitly:

CSV_COLUMNS = [
    "survived", "sex", "age", "n_siblings_spouses",
    "parch", "fare", "class", "deck", "embark_town", "alone"
]

temp_dataset = get_dataset(train_file_path, column_names=CSV_COLUMNS)

tf.data also supports selecting only the columns of interest directly out of the CSV, avoiding the need to load everything into memory and drop columns afterward (as would be required with pandas):

SELECT_COLUMNS = ["survived", "class", "deck", "embark_town", "alone"]

temp_dataset = get_dataset(train_file_path, select_columns=SELECT_COLUMNS)
# The resulting batch only contains: class, deck, embark_town, alone (plus survived as the label)

A key advantage here is that CSV loading through tf.data can select relevant columns up front, rather than loading everything and discarding what is not needed (the typical pandas workflow).

Loading Data from NumPy and pandas

NumPy and pandas are both excellent for data manipulation. pandas in particular is useful in almost any workflow that deals with structured, tabular data in a traditional machine learning approach. NumPy and pandas are less commonly used directly for image or text data, given how that data is stored.

From NumPy:

import numpy as np
import tensorflow as tf

DATA_URL = "https://storage.googleapis.com/tensorflow/tf-keras-datasets/mnist.npz"
path = tf.keras.utils.get_file("mnist.npz", DATA_URL)

with np.load(path) as data:
    train_examples = data["x_train"]
    train_labels = data["y_train"]
    test_examples = data["x_test"]
    test_labels = data["y_test"]

print(type(train_examples))  # numpy.ndarray

# Convert the NumPy arrays into a tf.data.Dataset
train_dataset = tf.data.Dataset.from_tensor_slices((train_examples, train_labels))
test_dataset = tf.data.Dataset.from_tensor_slices((test_examples, test_labels))

print(type(train_dataset))  # TensorSliceDataset

for example, label in train_dataset.take(1):
    print(example, label)   # values are tf.Tensor objects, similar in shape to the NumPy arrays

From pandas:

import pandas as pd
import tensorflow as tf

csv_file = tf.keras.utils.get_file(
    "heart.csv",
    "https://storage.googleapis.com/applied-dl/heart.csv"
)

df = pd.read_csv(csv_file)
print(df.shape)   # (303, 14)
print(df.head())
print(df.dtypes)  # mostly int, one float (oldpeak), one object/string (thal)

# Convert the string "thal" column into a categorical type
df["thal"] = pd.Categorical(df["thal"])
df["thal"] = df.thal.cat.codes  # keep the encoding but as numeric codes

print(df.dtypes)  # thal is now an integer

target = df.pop("target")

dataset = tf.data.Dataset.from_tensor_slices((df.values, target.values))
print(type(dataset))  # TensorSliceDataset

Note that df.values and target.values convert the pandas objects into NumPy arrays before being passed into from_tensor_slices — pandas data is effectively broken down to NumPy under the hood to build the tensor-slice dataset.

Serializing Data with tf.Example

tf.Example is a somewhat more complex data storage format, but it enables efficient storage and loading by allowing data to be converted into binaries. tf.Example holds data in memory as key/value pairs (similar to a dictionary), where each key maps to a typed list (e.g. int64_list, float_list, bytes_list). Multiple key/value pairs can be nested together to represent an entire record, and the whole structure can then be serialized into a compact binary string.

tf.Example is not strictly required before moving on to TFRecord, but it is highly recommended, since it makes serialization far more manageable. There is no need to convert existing code to use tf.Example/TFRecord unless tf.data is already in use and reading data remains the training bottleneck. If data fits comfortably in memory, it typically isn’t worth the additional effort.

import numpy as np
import pandas as pd
import tensorflow as tf

n_observations = 10000

# Four example feature vectors
feature0 = np.random.choice([False, True], n_observations)             # Boolean feature
feature1 = np.random.randint(0, 5, n_observations)                     # Integer feature
strings = np.array(["a", "b", "c", "d", "e"])
feature2 = strings[feature1]                                           # String feature tied to feature1
feature3 = np.random.randn(n_observations)                             # Float feature

features_dataset = tf.data.Dataset.from_tensor_slices((feature0, feature1, feature2, feature3))

for f0, f1, f2, f3 in features_dataset.take(1):
    print(f0, f1, f2, f3)
    # dtypes: bool, int, string, float -- matching the original NumPy arrays

# Helper functions (from the TensorFlow docs) to convert values into tf.train.Feature types
def _bytes_feature(value):
    if isinstance(value, type(tf.constant(0))):
        value = value.numpy()
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))

def _float_feature(value):
    return tf.train.Feature(float_list=tf.train.FloatList(value=[value]))

def _int64_feature(value):
    return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))

def serialize_example(feature0, feature1, feature2, feature3):
    feature = {
        "feature0": _int64_feature(feature0),
        "feature1": _int64_feature(feature1),
        "feature2": _bytes_feature(feature2),
        "feature3": _float_feature(feature3),
    }
    example_proto = tf.train.Example(features=tf.train.Features(feature=feature))
    return example_proto.SerializeToString()

# Serialize a single observation
serialized_example = serialize_example(False, 4, b"goat", 0.1234)
print(serialized_example)  # a binary string, not human-readable

# Deserialize it back into an Example object
example_proto = tf.train.Example.FromString(serialized_example)
print(example_proto)       # shows the feature dictionary, e.g. float_list { value: 0.1234 }
print(type(example_proto)) # tensorflow.core.example.example_pb2.Example

Storing Data Efficiently with TFRecord

TFRecord builds on tf.Example to enable easier, faster loading of data into pipelines. It is a simple binary format that compresses and decompresses efficiently, which matters when streaming data over a network or reading it back off disk quickly. TFRecord becomes important when a dataset does not fit entirely in memory, since data must then be loaded in batches as fast as possible.

TFRecord adds complexity, so it is not necessary when data comfortably fits in memory (e.g. typical structured data in a traditional ML workflow). It becomes valuable for large volumes of text, image, or other computer-vision-style data, particularly when the bottleneck is loading data off disk rather than training itself — loading from CSV is considerably slower than loading from a TFRecord binary.

sequenceDiagram
    participant NP as NumPy feature arrays
    participant FD as features_dataset
    participant SE as serialize_example()
    participant SD as serialized_dataset
    participant Disk as tf_data_pipelines.tfrecord
    participant RD as raw_dataset (TFRecordDataset)

    NP->>FD: from_tensor_slices((f0, f1, f2, f3))
    FD->>SE: map tf_serialize_example over each row
    SE->>SD: SerializeToString() per observation
    SD->>Disk: TFRecordWriter.write(serialized_dataset)
    Disk->>RD: TFRecordDataset(filenames)
    RD-->>RD: raw dataset in memory, ready to parse

import tensorflow as tf

def tf_serialize_example(f0, f1, f2, f3):
    tf_string = tf.py_function(
        serialize_example,
        (f0, f1, f2, f3),
        tf.string
    )
    return tf.reshape(tf_string, ())

# Map the serialization function across the entire features_dataset
serialized_features_dataset = features_dataset.map(tf_serialize_example)

# Alternative: build the same serialized dataset via a generator
def generator():
    for features in features_dataset:
        yield serialize_example(*features)

serialized_features_dataset = tf.data.Dataset.from_generator(
    generator, output_types=tf.string, output_shapes=()
)

# Write the serialized dataset out to disk as a .tfrecord file
filename = "tf_data_pipelines.tfrecord"
writer = tf.data.experimental.TFRecordWriter(filename)
writer.write(serialized_features_dataset)

# Read it back
filenames = [filename]
raw_dataset = tf.data.TFRecordDataset(filenames)

TFRecord requires a little more setup than keeping everything as a plain tf.data object, but it enables loading and unloading data from memory efficiently when working with larger datasets.

Loading Image Data

Beyond structured/tabular data, one of the most common uses of TensorFlow is working with image data. There are two general ways of importing image data: from a prepared set of data already available, or from a compressed archive (ZIP or TGZ). This section covers the prepared-data case; the ZIP/generator case is covered in Module 3.

import matplotlib.pyplot as plt
import tensorflow as tf

fashion_mnist = tf.keras.datasets.fashion_mnist
(train_images, train_labels), (test_images, test_labels) = fashion_mnist.load_data()

print(train_images.shape)  # (60000, 28, 28) -- a 3D tensor: observations x height x width

def plot_image(image):
    plt.figure()
    plt.imshow(image)
    plt.colorbar()
    plt.grid(False)
    plt.show()

plot_image(train_images[0])
plot_image(train_images[1])  # e.g. a T-shirt

# Pixel values range 0-255 -- scale down to 0-1 for the neural network
train_images = train_images / 255.0
test_images = test_images / 255.0  # the test set must receive the exact same transformation

Pixel values represent color intensity (0 is near-black, 255 is near-white/bright). Scaling every feature to the 0-1 range prevents columns with larger raw ranges from dominating the model purely due to scale. Critically, whatever transformation is applied to the training set must be applied identically to the test set — otherwise the model’s performance on unseen data will not reflect what was learned during training.

Module 3: Prepping Data

This module covers feature engineering and how to prep data before feeding it into machine learning models. Data preparation is widely considered one of the most difficult and time-consuming steps in machine learning — in practice, most of the effort in a real project goes into specifying the problem and getting data into the right shape, not into the modeling step itself. This module covers splitting data appropriately, creating additional features for structured data, and augmenting datasets when working with images.

Feature Engineering with pandas

This section is a light overview of pandas-based feature engineering: reading data, inspecting for missing values, converting data types, one-hot encoding, and splitting into train/test sets to guard against overfitting.

One-hot encoding vs. dummy variables:

Train/test split concepts:

• Train set — The data actually fed into the model during training (commonly ~80% of the full dataset).
• Test / validation / holdout set — The remaining data, held back from training and used only to evaluate performance metrics. If these observations leaked into training, the model could overfit and the reported performance would not reflect how it behaves on genuinely new data.
• Once a model’s performance looks acceptable on the holdout set, the entire dataset (train + test) is typically used for the final production model — the train/test split exists purely for evaluation, not for the final training run.

Demo: preparing the Auto MPG dataset

import numpy as np
import pandas as pd
import tensorflow as tf

column_names = [
    "MPG", "Cylinders", "Displacement", "Horsepower", "Weight",
    "Acceleration", "Model Year", "Origin"
]

raw_dataset = pd.read_csv(
    "auto-mpg.data",
    names=column_names,
    na_values="?",
    comment="\t",
    sep=" ",
    skipinitialspace=True
)

df = raw_dataset.copy()  # copy to avoid a SettingWithCopyWarning when modifying the data
print(df.head())
print(df.shape)          # (398, 8)
print(df.keys())         # or df.columns -- both list the column names

# Inspect missing values
print(df.isna().sum())   # Horsepower has 6 missing values, all other columns have 0

# Option 1: drop rows with any missing values
df1 = df.dropna()
print(df1.shape)         # (392, 8)

# Option 2: impute missing values with the column mean
df2 = df.fillna(df.mean())
print(df2.isna().sum())  # all 0s now

df = df2  # keep the imputed version to preserve more observations

print(df.dtypes)  # Origin shows up as an integer, but it actually encodes a country of origin

# Map the numeric origin codes to readable labels
df["Origin"] = df["Origin"].map({1: "USA", 2: "Europe", 3: "Japan"})
print(df.tail())

# One-hot encode the origin column
df = pd.get_dummies(df, prefix="", prefix_sep="")
print(df.head())  # now has separate Europe, Japan, USA columns

# Train/test split (80/20)
train_dataset = df.sample(frac=0.8, random_state=0)
test_dataset = df.drop(train_dataset.index)

# Pop the label column out of both sets
train_labels = train_dataset.pop("MPG")
test_labels = test_dataset.pop("MPG")

# Build tf.data datasets directly from the pandas data
train_data = tf.data.Dataset.from_tensor_slices((train_dataset.values, train_labels.values))
test_data = tf.data.Dataset.from_tensor_slices((test_dataset.values, test_labels.values))

Using tf.data is not mandatory when working from pandas, but converting to tensors is necessary to unlock batching, prefetching, and the other performance optimizations covered in Module 4.

Combining Datasets with zip and map

zip and map allow multiple tf.data objects to be combined and transformed into a single dataset.

import tensorflow as tf

dataset1 = tf.data.Dataset.from_tensor_slices(tf.random.uniform([4, 10]))
print(dataset1.element_spec)  # TensorSpec(shape=(10,), dtype=tf.float32)

dataset2 = tf.data.Dataset.from_tensor_slices(
    (tf.random.uniform([4]), tf.random.uniform([4, 100]))
)
print(dataset2.element_spec)
# (TensorSpec(shape=(), dtype=tf.float32), TensorSpec(shape=(100,), dtype=tf.float32))

# Combine dataset1 and dataset2 element-wise
dataset3 = tf.data.Dataset.zip((dataset1, dataset2))
print(dataset3.element_spec)
# Three TensorSpecs stacked together, one from dataset1 and two from dataset2

# map and filter
dataset = tf.data.Dataset.range(100)
list(dataset.as_numpy_iterator())  # values 0-99

mapped = dataset.map(lambda x: [x])
padded = mapped.padded_batch(4, padded_shapes=(None,))

padded_batch is used here to batch tensors of potentially different sizes together — this matters when moving toward GPU training, since it helps ensure batches stay consistent and fit into memory correctly.

Loading Image Data with Generators

Image data requires additional preprocessing steps compared to tabular data:

flowchart LR
    A["Read image bytes\nfrom disk"] --> B["Decode image\n(2D -> RGB channels)"]
    B --> C["Convert to tensor\n(tf.data object)"]
    C --> D["Scale pixel values\n0-255 -> 0-1"]

The final step before fitting a model is to scale images so that pixel values fall between 0 and 1 (or -1 to 1 for other data types) — image pixel values are always positive, and scaling avoids the model implicitly weighting one channel/column more heavily purely due to range.

Demo: cats vs. dogs with ImageDataGenerator

import os
import matplotlib.pyplot as plt
import tensorflow as tf
from tensorflow.keras.preprocessing.image import ImageDataGenerator

_URL = "https://storage.googleapis.com/mledu-datasets/cats_and_dogs_filtered.zip"
path_to_zip = tf.keras.utils.get_file("cats_and_dogs.zip", origin=_URL, extract=True)
base_dir = os.path.join(os.path.dirname(path_to_zip), "cats_and_dogs_filtered")

train_dir = os.path.join(base_dir, "train")
train_cats_dir = os.path.join(train_dir, "cats")

batch_size = 128
IMG_HEIGHT = 150
IMG_WIDTH = 150

image_generator = ImageDataGenerator(rescale=1.0 / 255)

train_data_gen = image_generator.flow_from_directory(
    batch_size=batch_size,
    directory=train_dir,
    shuffle=True,
    target_size=(IMG_HEIGHT, IMG_WIDTH)
)

def plot_image(image):
    plt.figure()
    plt.imshow(image)
    plt.colorbar()
    plt.grid(False)
    plt.show()

# First image from the generator (a cat)
plot_image(train_data_gen[0][0][0])
# Second image from the generator (a dog)
plot_image(train_data_gen[0][0][1])

def plot_5_images(images):
    fig, axes = plt.subplots(1, 5, figsize=(20, 20))
    axes = axes.flatten()
    for img, ax in zip(images, axes):
        ax.imshow(img)
    plt.tight_layout()
    plt.show()

images = [train_data_gen[0][0][i] for i in range(5)]
plot_5_images(images)

Generators like ImageDataGenerator.flow_from_directory offer performance benefits over loading everything into memory up front: images are streamed from disk as needed and do not persist in memory once consumed.

Image Data Augmentation

Neural networks can overfit to incidental characteristics of a small image set — for example, the exact spatial relationship between an eye and a nose in a centered photo. If a new image shows the same subject rotated, shifted, or viewed from a different angle, a human can still recognize it instantly, but the model may fail to generalize. Image data augmentation synthetically creates additional training examples from each original image to counter this.

mindmap
  root((Image Data Augmentation))
    rotation_range
      "Rotate image up to N degrees (0-360)"
    width_shift_range
      "Squeeze/shift image left-right"
    height_shift_range
      "Squeeze/shift image up-down"
    horizontal_flip
      "Mirror the image"
    zoom_range
      "Zoom in/out on a portion of the image"

from tensorflow.keras.preprocessing.image import ImageDataGenerator

# Horizontal flip only
image_gen = ImageDataGenerator(rescale=1.0 / 255, horizontal_flip=True)
train_data_gen = image_gen.flow_from_directory(
    batch_size=batch_size, directory=train_dir,
    shuffle=True, target_size=(IMG_HEIGHT, IMG_WIDTH)
)

# Rotation only
image_gen = ImageDataGenerator(rescale=1.0 / 255, rotation_range=45)
train_data_gen = image_gen.flow_from_directory(
    batch_size=batch_size, directory=train_dir,
    shuffle=True, target_size=(IMG_HEIGHT, IMG_WIDTH)
)

# Zoom only
image_gen = ImageDataGenerator(rescale=1.0 / 255, zoom_range=0.5)
train_data_gen = image_gen.flow_from_directory(
    batch_size=batch_size, directory=train_dir,
    shuffle=True, target_size=(IMG_HEIGHT, IMG_WIDTH)
)

# Combine all augmentations together
image_gen_train = ImageDataGenerator(
    rescale=1.0 / 255,
    rotation_range=45,
    width_shift_range=0.15,
    height_shift_range=0.15,
    horizontal_flip=True,
    zoom_range=0.5
)

train_data_gen = image_gen_train.flow_from_directory(
    batch_size=batch_size, directory=train_dir,
    shuffle=True, target_size=(IMG_HEIGHT, IMG_WIDTH)
)

images = [train_data_gen[0][0][0] for _ in range(5)]
plot_5_images(images)
# Each of the 5 images now differs: some zoomed, some rotated,
# some flipped, some shifted -- all derived from the same source photo

Combining rotation, shifting, flipping, and zooming yields a large number of unique variations from a single source image. This is especially valuable given that one of the most expensive parts of image data prep is obtaining and labeling images in the first place — augmentation effectively multiplies the available training data for free.

Module 4: Optimizing Performance of Pipelines

This module covers how to squeeze performance out of tf.data datasets. Accelerating computation with a GPU or TPU can dramatically lower training time — but if data does not show up at the right time, loading the data itself becomes the new bottleneck, especially with modern datasets too large to fit in memory. This module reviews the Keras API at a high level (needed context for reasoning about pipeline optimizations), how to evaluate model performance, and then dives into the two central optimizations: batching, prefetching, and finally parallelizing data extraction.

flowchart LR
    A["Raw data\n(pandas / CSV / NumPy)"] --> B["tf.data.Dataset\n.from_tensor_slices"]
    B --> C["batch(n)\nsplit into fixed-size batches"]
    C --> D["prefetch(n)\nload next batch while training"]
    D --> E["interleave + AUTOTUNE\nparallelize remote/disk reads"]
    E --> F["model.fit(dataset, epochs=...)"]

Preparing Data for Model Training

This section revisits the heart disease dataset (introduced in Module 2) to walk through normalization, a prerequisite step for effective neural network training.

import time
import pandas as pd
import tensorflow as tf
from tensorflow import keras

df = pd.read_csv("heart.csv")
print(df.head())  # 14 columns, including "target" (the label to predict)

df["thal"] = pd.Categorical(df["thal"])
df["thal"] = df.thal.cat.codes

train_dataset = df.sample(frac=0.8, random_state=0)
test_dataset = df.drop(train_dataset.index)

train_labels = train_dataset.pop("target")
test_labels = test_dataset.pop("target")

train_stats = train_dataset.describe().transpose()

def norm(x, stats):
    return (x - stats["mean"]) / stats["std"]

normed_train_data = norm(train_dataset, train_stats)
# Crucially, the *training* statistics (mean/std) are reused to normalize the test set --
# the model must never "see" statistics derived from data it hasn't been trained on.
normed_test_data = norm(test_dataset, train_stats)

Normalizing scales every column to a comparable range (commonly 0-1 or -1 to 1), which prevents columns with naturally larger raw values from dominating the network purely due to scale. The same normalization parameters (mean and standard deviation) computed from the training set must be reused — never recomputed — on the test set.

Building Models with the Keras Sequential API

model = keras.Sequential([
    keras.layers.Dense(64, activation="relu", input_shape=[len(normed_train_data.keys())]),
    keras.layers.Dense(64, activation="relu"),
    keras.layers.Dense(1)
])

model.compile(
    optimizer="rmsprop",
    loss="mse",
    metrics=["mse", "mae"]
)

model.fit(normed_train_data, train_labels, epochs=10)
# Loss/MSE/MAE decrease across epochs

y_hats = model.predict(normed_test_data)
y_hats = [item for sublist in y_hats for item in sublist]  # flatten predictions into a 1D array

from sklearn.metrics import mean_absolute_error
mean_absolute_error(test_labels, y_hats)  # e.g. 0.26536

# TensorFlow's built-in evaluate() achieves the same result more directly
results = model.evaluate(test_dataset, test_labels)
print(results)  # [loss, mse, mae] -- three values because two metrics were specified

Layer notes:

• A single output node is used here because this is a regression task (a single predicted value). A classification task with 10 classes would instead use a final Dense(10) layer.
• For image inputs, a Flatten layer is typically the first layer, converting 2D image data into a 1D shape before the dense layers.
• compile() requires three things: an optimizer (rmsprop is typical for this style of tabular data, though it is a good candidate for a grid search), a loss function, and one or more metrics.
• model.evaluate(x, y) runs prediction and metric computation in one call — the last value returned corresponds to the last metric specified (here, MAE), matching a manually computed MAE via scikit-learn.

Batching and Prefetching

Batching splits a dataset into smaller pieces; prefetching starts reading the next batch off disk before the model actually needs it, which helps remove data loading as a training bottleneck.

import tensorflow as tf

# Build a dataset directly from the normalized pandas data
dataset = tf.data.Dataset.from_tensor_slices((normed_train_data.values, train_labels.values))

# Without batching, model.fit will not accept the raw tuple-shaped dataset as-is.
# Batching resolves this by grouping observations into fixed-size chunks.
dataset = dataset.batch(100)

model.fit(dataset, epochs=100)
# Epoch output shows "3/3" -- roughly 240+ observations split into 3 batches of 100

dataset_small_batch = tf.data.Dataset.from_tensor_slices(
    (normed_train_data.values, train_labels.values)
).batch(10)

model.fit(dataset_small_batch, epochs=100)
# Epoch output shows "25/25" -- same ~242 observations now split into 25 batches of 10
# (training time increases slightly because more, smaller batches must each be loaded)

# Prefetching: start loading the next batch while the current one is still training on
dataset = dataset.prefetch(2)  # keep 2 batches ready ahead of time
model.fit(dataset, epochs=100)

Without prefetching, the model waits idle for each batch to be fetched and loaded before continuing; with prefetch(n), n batches are proactively loaded into memory so they are ready the instant the current batch finishes training. Batching and prefetching are recommended on essentially every dataset — there is little downside, and the win is substantial whenever loading is not already instantaneous.

Parallelizing Data Extraction with interleave

Even with batching and prefetching, getting data from disk (or especially from remote storage) into the training loop can remain the bottleneck. This matters most when data lives in a remote location — for example, on AWS S3 or on a server geographically far from the compute (imagine compute on the US west coast and storage on the US east coast, or even farther away in Europe). Two factors dominate here:

• Time to first byte — The latency before any data starts arriving from remote storage can be an order of magnitude larger than reading from local disk.
• Read throughput — Remote storage may offer significant available bandwidth, but reading sequentially fails to use all of it. Parallelizing reads allows more of that available bandwidth to be used simultaneously.

interleave is the tf.data transformation used to parallelize data loading:

import tensorflow as tf

dataset = tf.data.Dataset.from_tensor_slices([...])  # a tf.data.Dataset of file paths, for example

# Naive interleave -- defaults to sequential loading, no performance benefit on its own
dataset = dataset.interleave(lambda x: tf.data.Dataset.from_tensors(x))

# Parallel interleave with AUTOTUNE -- lets TensorFlow choose the level of parallelism
dataset = dataset.interleave(
    lambda x: tf.data.Dataset.from_tensors(x),
    num_parallel_calls=tf.data.experimental.AUTOTUNE
)

interleave accepts a cycle_length (how many source datasets/files to interleave from concurrently) and a level of parallelism (how many threads to open for loading). By default, interleave loads two datasets sequentially and offers no speedup — the num_parallel_calls argument is what actually enables parallel loading. Using tf.data.experimental.AUTOTUNE lets TensorFlow automatically decide the right degree of parallelism at runtime based on available threads and data, rather than hardcoding a thread count that risks thread-locking a machine running other processes.

flowchart TB
    A["Remote/disk file paths"] --> B["interleave(...)"]
    B --> C{"num_parallel_calls\nspecified?"}
    C -->|"No (default)"| D["Sequential loading\nNo performance benefit"]
    C -->|"AUTOTUNE"| E["TensorFlow picks parallelism\nbased on available threads/data"]
    E --> F["Multiple files/batches\nloaded concurrently"]

Best practices for optimizing a pipeline where data loading (not model training) is the bottleneck:

batch and prefetch should be used on essentially every tf.data pipeline, with little downside beyond very small datasets where it simply doesn’t matter either way. interleave becomes relevant specifically when large data volumes and/or remote storage make loading itself the constraint.

Summary

tf.data is one of the most impactful additions in TensorFlow 2.0, turning data loading from a hidden bottleneck into a tunable, optimizable stage of the training pipeline. Key takeaways:

• TensorFlow 2.0 consolidated the Keras API, added eager execution by default, introduced the TFX framework for production deployment, and — most importantly for this course — introduced tf.data for efficient, flexible data pipelines. Full backward compatibility with TensorFlow 1.0 code is possible via tf.compat.v1.disable_v2_behavior(), but doing so forfeits access to tf.data.
• tf.data.Dataset can be constructed either directly from an existing data source or by transforming other formats (CSV, NumPy, pandas, images) into a dataset via from_tensor_slices or the CSV-specific utilities.
• tf.Example and TFRecord provide an efficient binary serialization format for large datasets, particularly valuable for image/text data or any case where reading from disk is the bottleneck — but they add complexity that isn’t worth it for smaller, in-memory-friendly datasets.
• Feature engineering with pandas (missing value handling, one-hot encoding/categorical codes, train/test splitting) remains a core part of preparing structured data, even when the final data is converted into a tf.data.Dataset for training.
• Image data requires an additional pipeline: reading from disk, decoding into RGB tensors, and scaling pixel values to 0-1. ImageDataGenerator streams images efficiently and supports built-in data augmentation (rotation, shifting, flipping, zooming) to multiply a small labeled image set without collecting new data.
• Performance optimization boils down to three composable techniques: batch() to control memory usage and epoch granularity, prefetch() to overlap data loading with training computation, and interleave(..., num_parallel_calls=tf.data.experimental.AUTOTUNE) to parallelize reads when disk or remote-storage throughput is the true bottleneck.

Quick-Reference Checklist

• Decide whether existing TensorFlow 1.0 code needs true migration, or whether tf.compat.v1.disable_v2_behavior() is an acceptable (but limiting) shortcut.
• Load raw data into a tf.data.Dataset using the appropriate method: make_csv_dataset for CSV, from_tensor_slices for NumPy/pandas, ImageDataGenerator.flow_from_directory for images.
• Convert to tf.Example/TFRecord only if data no longer fits comfortably in memory or disk I/O is a measured bottleneck.
• Handle missing values, encode categorical columns (one-hot or categorical codes), and split into train/test sets before training.
• Normalize/scale numeric features and image pixel values using statistics derived only from the training set.
• Apply image augmentation (rotation_range, width_shift_range, height_shift_range, horizontal_flip, zoom_range) when working with a limited image dataset.
• Always apply batch() and prefetch() to training datasets.
• Apply interleave() with num_parallel_calls=tf.data.experimental.AUTOTUNE when data is remote or the dataset is too large to load sequentially fast enough.

Search Terms

designing · data · pipelines · tensorflow · deep · neural · networks · machine · science · loading · image · 2.0 · model · pandas