Building & Scaling Models with TensorFlow

Components Overview

TensorFlow is far more than a single library — it is an end-to-end platform for building, training, and deploying machine learning models. Understanding its ecosystem is essential for the MLE exam because many questions test whether you know which component to reach for in a given scenario.

Reading Data from Various Sources

The tf.data.Dataset API is TensorFlow's standard mechanism for building efficient, scalable input pipelines. It provides a functional, chainable API that handles reading, parsing, transforming, batching, and prefetching data. On GCP, your data typically lives in BigQuery, Cloud Storage (GCS), or as TFRecord files.

TFRecord format is TensorFlow's native binary serialization format. It stores data as a sequence of serialized protocol buffers (tf.train.Example), each containing a dictionary of features. TFRecords are ideal for large datasets because they support efficient sequential reads, work well with GCS, and can be sharded across multiple files for parallel reading.

# Reading from GCS (CSV files)
dataset = tf.data.experimental.make_csv_dataset(
    "gs://my-bucket/data/train-*.csv",
    batch_size=32,
    label_name="target",
    num_epochs=1,
    shuffle=True
)

# Reading TFRecord files from GCS
filenames = tf.io.gfile.glob("gs://my-bucket/data/train-*.tfrecord")
dataset = tf.data.TFRecordDataset(filenames, compression_type="GZIP")
dataset = dataset.map(parse_fn, num_parallel_calls=tf.data.AUTOTUNE)

# Reading from BigQuery via BigQuery Connector
from tensorflow_io.bigquery import BigQueryClient
client = BigQueryClient()
read_session = client.read_session(
    "projects/my-project",
    parent_project="my-project",
    table="my-project.dataset.table",
    selected_fields=["feature1", "feature2", "label"],
    output_types=[tf.float32, tf.float32, tf.int32],
    requested_streams=4
)
dataset = read_session.parallel_read_rows()

tf.data Performance Optimization

A naively constructed input pipeline will bottleneck your training. The GPU sits idle while waiting for data to load, parse, and transform. The tf.data API provides several mechanisms to overlap I/O with computation and maximize throughput.

Prefetch overlaps data preprocessing with model training. While the GPU processes batch N, the CPU prepares batch N+1. Always use tf.data.AUTOTUNE to let TensorFlow dynamically tune the buffer size.

Interleave reads from multiple files simultaneously, reducing I/O bottlenecks when data is sharded across many files. The cycle_length parameter controls how many files are read concurrently.

Cache stores the dataset in memory (or on local disk) after the first epoch, eliminating redundant I/O for subsequent epochs. Use .cache() after expensive transformations but before randomized augmentations.

Parallel map uses multiple CPU cores to apply transformation functions. Set num_parallel_calls=tf.data.AUTOTUNE for automatic parallelism tuning.

# Optimized input pipeline pattern
def build_dataset(file_pattern, batch_size=64):
    files = tf.data.Dataset.list_files(file_pattern, shuffle=True)

    dataset = files.interleave(
        lambda f: tf.data.TFRecordDataset(f),
        cycle_length=8,              # Read from 8 files concurrently
        num_parallel_calls=tf.data.AUTOTUNE,
        deterministic=False          # Non-deterministic for speed
    )

    dataset = dataset.map(
        parse_and_augment,
        num_parallel_calls=tf.data.AUTOTUNE
    )

    dataset = dataset.cache()             # Cache after parsing
    dataset = dataset.shuffle(10000)       # Shuffle buffer
    dataset = dataset.batch(batch_size)
    dataset = dataset.prefetch(tf.data.AUTOTUNE)  # Overlap with training

    return dataset

Sequential and Functional APIs

Keras provides three progressively flexible ways to define models. The Sequential API is the simplest: it creates a linear stack of layers. It works for straightforward feed-forward networks but cannot handle multiple inputs, multiple outputs, shared layers, or non-linear topologies.

# Sequential API — linear stack of layers
model = tf.keras.Sequential([
    tf.keras.layers.Dense(128, activation="relu", input_shape=(784,)),
    tf.keras.layers.Dropout(0.3),
    tf.keras.layers.Dense(64, activation="relu"),
    tf.keras.layers.Dense(10, activation="softmax")
])

The Functional API supports arbitrary graph topologies. You define layers as functions called on tensors, enabling multi-input/multi-output architectures, residual connections, and shared layers. This is the most commonly used API for production models.

# Functional API — arbitrary graph topologies
inputs = tf.keras.Input(shape=(784,))
x = tf.keras.layers.Dense(128, activation="relu")(inputs)
x = tf.keras.layers.Dropout(0.3)(x)
residual = x
x = tf.keras.layers.Dense(128, activation="relu")(x)
x = tf.keras.layers.Add()([x, residual])    # Residual connection
outputs = tf.keras.layers.Dense(10, activation="softmax")(x)

model = tf.keras.Model(inputs=inputs, outputs=outputs)

Model Subclassing

Model Subclassing gives you full control by defining the model as a Python class. You implement __init__ to create layers and call to define the forward pass. This enables dynamic architectures with Python control flow (loops, conditionals) but sacrifices some Keras features like automatic serialization and layer visualization.

# Model Subclassing — full Python flexibility
class CustomModel(tf.keras.Model):
    def __init__(self, num_classes=10):
        super().__init__()
        self.dense1 = tf.keras.layers.Dense(128, activation="relu")
        self.dropout = tf.keras.layers.Dropout(0.3)
        self.dense2 = tf.keras.layers.Dense(num_classes, activation="softmax")

    def call(self, inputs, training=False):
        x = self.dense1(inputs)
        x = self.dropout(x, training=training)  # Dropout active only during training
        return self.dense2(x)

Architecture Selection Guide

Choosing the right architecture depends on the data type and the task. The exam tests your ability to match architectures to problems.

Keras callbacks are objects that hook into the training loop at specific points (epoch start/end, batch start/end) to perform actions like saving checkpoints, adjusting learning rates, or stopping training early. They are essential for production training workflows.

# Essential callbacks for production training
callbacks = [
    # Stop training when validation loss stops improving
    tf.keras.callbacks.EarlyStopping(
        monitor="val_loss",
        patience=5,            # Wait 5 epochs for improvement
        restore_best_weights=True,  # Revert to best epoch
        min_delta=0.001        # Minimum change to qualify
    ),

    # Save the best model checkpoint
    tf.keras.callbacks.ModelCheckpoint(
        filepath="best_model.keras",
        monitor="val_loss",
        save_best_only=True,
        save_weights_only=False  # Save full model
    ),

    # Log metrics for TensorBoard visualization
    tf.keras.callbacks.TensorBoard(
        log_dir="./logs",
        histogram_freq=1,    # Log weight histograms
        profile_batch="10,20" # Profile batches 10-20
    ),

    # Reduce learning rate when metric plateaus
    tf.keras.callbacks.ReduceLROnPlateau(
        monitor="val_loss",
        factor=0.5,            # Halve the learning rate
        patience=3,            # Wait 3 epochs
        min_lr=1e-6            # Minimum learning rate
    ),
]

# Use callbacks in model.fit()
history = model.fit(
    train_dataset,
    validation_data=val_dataset,
    epochs=100,             # Set high; EarlyStopping will halt
    callbacks=callbacks
)

While model.fit() handles most training scenarios, you sometimes need fine-grained control over the training step — for example, when implementing GANs, multi-task learning with custom loss weighting, or gradient accumulation for large effective batch sizes. TensorFlow's tf.GradientTape records operations for automatic differentiation.

# Custom training loop with tf.GradientTape
optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3)
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy()

for epoch in range(num_epochs):
    for x_batch, y_batch in train_dataset:
        with tf.GradientTape() as tape:
            predictions = model(x_batch, training=True)
            loss = loss_fn(y_batch, predictions)

        # Compute gradients
        gradients = tape.gradient(loss, model.trainable_variables)

        # Optional: clip gradients to prevent exploding gradients
        gradients = [tf.clip_by_norm(g, 1.0) for g in gradients]

        # Apply gradients
        optimizer.apply_gradients(zip(gradients, model.trainable_variables))

When to use custom training loops: GANs (alternating generator/discriminator updates), reinforcement learning, gradient accumulation for large batch simulation, multi-task learning with independent optimizers, or when you need to log intermediate values that model.fit() does not expose.

# Optimized custom step with @tf.function
@tf.function
def train_step(x, y):
    with tf.GradientTape() as tape:
        predictions = model(x, training=True)
        loss = loss_fn(y, predictions)
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    return loss

Distribution Strategies

TensorFlow's tf.distribute module provides strategies for distributing training across multiple devices and machines. The key insight is that most strategies implement data parallelism: the model is replicated on each device, each replica processes a slice of the data, and gradients are aggregated (typically via all-reduce) before updating.

# MirroredStrategy — single machine, multi-GPU
strategy = tf.distribute.MirroredStrategy()
print(f"Number of devices: {strategy.num_replicas_in_sync}")

with strategy.scope():
    model = build_model()
    model.compile(
        optimizer="adam",
        loss="sparse_categorical_crossentropy",
        metrics=["accuracy"]
    )

# Global batch size = per_replica_batch * num_replicas
GLOBAL_BATCH = 64 * strategy.num_replicas_in_sync
train_ds = build_dataset(batch_size=GLOBAL_BATCH)

model.fit(train_ds, epochs=10)

# TPUStrategy
resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)
tpu_strategy = tf.distribute.TPUStrategy(resolver)

Custom Training Jobs & Hyperparameter Tuning

Vertex AI provides managed infrastructure for training ML models at any scale. You can submit custom training jobs that run your training code on Google Cloud compute resources, with options for pre-built containers (TensorFlow, PyTorch, XGBoost) or custom Docker containers.

Pre-built containers are Google-maintained Docker images with popular ML frameworks pre-installed. Use them when your training code uses standard frameworks and dependencies. They save time and reduce container management overhead.

Custom containers give you full control over the runtime environment. Use them when you have custom dependencies, non-standard frameworks, or need specific system libraries. You build a Docker image, push it to Artifact Registry, and reference it in your training job.

# Vertex AI Custom Training Job with SDK
from google.cloud import aiplatform

aiplatform.init(project="my-project", location="us-central1")

# Pre-built container approach
job = aiplatform.CustomTrainingJob(
    display_name="tf-training-job",
    script_path="trainer/task.py",
    container_uri="us-docker.pkg.dev/vertex-ai/training/tf-gpu.2-12:latest",
    requirements=["cloudml-hypertune"],
    model_serving_container_image_uri="us-docker.pkg.dev/vertex-ai/prediction/tf2-gpu.2-12:latest",
)

model = job.run(
    replica_count=1,
    machine_type="n1-standard-8",
    accelerator_type="NVIDIA_TESLA_V100",
    accelerator_count=2,
    args=["--epochs=50", "--batch-size=128"],
)

# Custom container approach
custom_job = aiplatform.CustomContainerTrainingJob(
    display_name="custom-container-job",
    container_uri="us-docker.pkg.dev/my-project/my-repo/trainer:latest",
    model_serving_container_image_uri="us-docker.pkg.dev/vertex-ai/prediction/tf2-gpu.2-12:latest",
)

Hyperparameter tuning on Vertex AI uses the Vizier service to efficiently search the hyperparameter space. You define the parameters to tune, their ranges, the optimization metric, and the search algorithm. Vertex AI manages trial scheduling, early stopping of unpromising trials, and resource allocation.

# Hyperparameter Tuning Job
from google.cloud.aiplatform import hyperparameter_tuning as hpt

job = aiplatform.HyperparameterTuningJob(
    display_name="hpt-tf-model",
    custom_job=custom_training_job,
    metric_spec={"val_accuracy": "maximize"},
    parameter_spec={
        "learning_rate": hpt.DoubleParameterSpec(
            min=1e-5, max=1e-1, scale="log"
        ),
        "num_layers": hpt.IntegerParameterSpec(
            min=2, max=8, scale="linear"
        ),
        "dropout_rate": hpt.DoubleParameterSpec(
            min=0.1, max=0.5, scale="linear"
        ),
    },
    max_trial_count=50,
    parallel_trial_count=5,
    search_algorithm=None,  # None = Bayesian (Vizier default)
)

job.run()

This section summarizes the key concepts from Course 07 that are most likely to appear on the GCP Professional Machine Learning Engineer exam, mapping to Section 3: Scaling prototypes into ML models.

Distribution Strategy Selection

• MirroredStrategy: single machine, multiple GPUs. Most common. Uses NCCL all-reduce.
• MultiWorkerMirroredStrategy: multiple machines. Requires TF_CONFIG. Use when single machine GPUs are insufficient.
• TPUStrategy: for TPU pods. Requires TPUClusterResolver. Prefer bfloat16 precision.
• Batch size must scale with replicas. Learning rate may need linear scaling + warmup.

tf.data Optimization

• Always use prefetch(tf.data.AUTOTUNE) at the end of your pipeline.
• Use interleave for parallel reading from multiple files.
• Cache after parsing, before augmentation. Shuffle after cache.
• TFRecord is the preferred format for large-scale training on GCP.

Container Selection

• Pre-built containers: standard frameworks, pip-installable dependencies.
• Custom containers: proprietary code, custom system libraries, specific CUDA versions.
• Hyperparameter tuning: use Vizier (Bayesian optimization), define metric_spec and parameter_spec.
• Use cloudml-hypertune to report metrics from training code to the tuning service.

How to Explain This in 2 Minutes

Likely Interview Questions

Model Answers

System Design Scenario

Common Mistakes

• Putting prefetch before batch or cache after shuffle — tf.data pipeline ordering is critical. The correct order is: read → parse/map → cache → shuffle → batch → prefetch. Caching after shuffle wastes the cache (different order each epoch). Prefetching before batching means you prefetch individual examples instead of full batches, losing the overlap benefit.
• Not scaling batch size and learning rate with distribution strategy — When using MirroredStrategy with 4 GPUs, the effective batch size is 4x the per-replica batch. If you do not increase the learning rate proportionally (with warmup), training will converge to a worse optimum because each gradient step averages over more examples without a correspondingly larger update. The standard recipe is: linear LR scaling with a warmup period of 5-10% of total steps.
• Using Model Subclassing when the Functional API would suffice — Subclassed models lose the benefits of Keras's graph tracing: no automatic shape inference, no model.summary(), harder serialization, and no easy export to SavedModel format. Unless your forward pass requires dynamic Python control flow, the Functional API is strictly preferable. Interviewers view unnecessary subclassing as a sign of inexperience with the Keras ecosystem.