Advanced TensorFlow: Custom Training and Optimization

Table of Contents

• Module 1: Custom Training Loops — The Production Baseline
  • When model.fit() Is Not Enough
  • Building the Complete GradientTape Loop
  • Correctness Conditions of a Custom Training Step
  • Demo: Train a Small Network End-to-End
• Module 2: Optimization Controls — Clipping and Schedules
  • Why Training Becomes Unstable
  • Implementing Gradient Clipping and LR Schedules
  • Tuning and Validating Stabilizers in Production
  • Demo: Stabilize Training with Clipping and Schedule
• Module 3: Advanced Loop Techniques — Mixed Precision and Nested Tape
  • Mixed Precision Mechanics That Matter
  • Nested GradientTape and Higher-Order Differentiation
  • Engineering Pattern: Safe Mixed Precision and Nested Tape
  • Demo: Mixed Precision with Gradient Penalty
• Module 4: Custom Layers and Models — Subclassing for Real Architectures
  • When to Subclass Layer and Model
  • Implementing Custom Components
  • Production Readiness: Tracking and Saving Subclassed Components
  • Demo: Plug Custom Components into Training
• Module 5: Custom Training Components — Losses, Metrics, Optimizers, and train_step()
  • Designing Custom Losses and Metrics
  • Custom Optimizers and train_step()
  • Validating Custom Components Before You Trust Them
  • Demo: Specialized Training with train_step()
• Module 6: Distributed Training — Strategies That Ship
  • Strategy Selection and Trade-offs
  • Correctness and Performance in Distributed Training
  • Converting Code to Distributed Execution
  • Demo: Distributed Custom Loop
• Module 7: Architecture Spotlights — GNNs, VAEs, GANs, RL, and scikit-learn Integration
  • Pattern Map for Advanced Workflows
  • Architecture-Driven Training Step Patterns
  • Building GNN Layers and Models
  • VAEs and GANs with Custom Loops
  • Demo: RL Loop and scikit-learn Integration

Module 1: Custom Training Loops — The Production Baseline

When model.fit() Is Not Enough

Modern production systems demand more than convenience APIs. Understanding what really happens inside a training step — how gradients are computed, how updates are applied, and how instability appears — enables structural reasoning about training, not just goal-oriented API usage.

There comes a point in advanced workflows where high-level orchestration becomes restrictive. Abstraction stops helping and starts hiding critical behavior. The shift is from convenience-driven training to execution-driven training.

Limits of High-Level Keras Training

Five concrete limitations emerge when model.fit() is used for complex workflows:

1. Training steps are abstracted away. The order of forward execution, loss reduction, gradient recording, and optimizer updates is handled internally. Direct intervention between those stages is not possible.

2. Gradient flow is not directly observable. If gradients explode, vanish, saturate, or produce NaNs, you typically discover through symptoms rather than through inspection at the source.

3. Loss composition becomes constrained. While multiple losses can be defined, dynamically assembling structured objectives or conditionally modifying losses per batch becomes cumbersome.

4. Debugging relies on indirect callbacks. Inspection happens outside the actual update logic. You cannot pause mid-step and inspect intermediate tensors.

5. Update logic is largely predefined. Non-standard flows — alternating optimizers, gradient accumulation, multi-phase updates, or conditional backpropagation — become difficult or awkward to express.

These constraints are not problems in simple workflows. They become problems when architectural complexity increases.

model.fit() vs. Custom Training Loop

graph LR
    A["model.fit()"] -->|Fixed flow| B[Predefined pipeline]
    C[Custom Loop] -->|Explicit control| D[Every stage written explicitly]
    A -->|Single optimizer| E[Uniform updates]
    C -->|Multiple/conditional| F[Separate optimizers per variable group]
    A -->|Limited loss| G[Predefined loss structure]
    C -->|Arbitrary composition| H[Dynamic multi-term objectives]
    A -->|Callback debugging| I[Post-update inspection]
    C -->|Inline inspection| J[Gradients, norms, tensors mid-step]

The distinction is not about replacing one with the other. It is about recognizing when abstraction is sufficient and when structural control becomes essential.

Building the Complete GradientTape Loop

Control requires understanding structure. A custom loop is not just more code — it is explicit control over execution. Instead of relying on abstractions, each step is made visible and intentional.

Training Loop Execution Flow

flowchart LR
    A[Input Batch] --> B[Forward Pass]
    B --> C[Loss Compute]
    C --> D[GradientTape Records]
    D --> E[Gradient Calculation]
    E --> F[Weight Update]

Training starts with the input batch, typically coming from a dataset pipeline. That batch flows into the forward pass, where the model produces predictions. Next, the loss is computed — it must reduce to a scalar value. Inside the GradientTape, TensorFlow records the operations performed on trainable variables during the forward pass. From that recorded graph, gradients are computed. Finally, the optimizer applies those gradients to update the weights. Nothing is implicit; each step must be defined deliberately.

Structural Correctness Rules

1. The forward pass must execute inside the GradientTape context. If it happens outside, TensorFlow records nothing and gradients will not exist.

2. The loss must reduce to a single scalar value. Differentiation operates on scalar objectives. Without proper reduction, gradient scaling becomes incorrect or the computation can break entirely.

3. Gradient calculation must produce exactly one tensor per trainable variable. A None gradient signals that a variable was disconnected from the computation graph or never participated in the forward pass.

4. The optimizer must apply updates to exactly the same variable list used during gradient computation. Inconsistent filtering or altering introduces subtle training bugs.

Minimal Reusable Loop Structure

The essential components of a production-ready custom loop:

flowchart TD
    A[Dataset Iteration] --> B[GradientTape Scope]
    B --> C[Explicit Loss Computation]
    C --> D[Manual Gradient Calculation]
    D --> E[Controlled Parameter Updates]
    E --> F[Optional Logging and Checks]

• Dataset iteration — Handle iteration explicitly to validate batch structure, confirm data alignment, and reason about how input flows.
• GradientTape scope — Defines what TensorFlow records for differentiation. The forward pass and all differentiable computations must happen here.
• Loss computation — Assembled intentionally, combining the main objective with penalties or regularization, reduced to a single scalar.
• Manual gradient calculation — Exposes raw gradients for inspection, validation, or transformation before any parameters are updated.
• Controlled parameter updates — The optimizer applies updates. Explicit control allows conditional updates, clipping, or multiple optimizers.
• Logging and checks — Tracking metrics, gradient norms, or numerical stability makes the loop easier to debug and safer for production.

A concrete minimal training step:

import tensorflow as tf

@tf.function
def train_step(model, optimizer, loss_fn, x_batch, y_batch):
    with tf.GradientTape() as tape:
        predictions = model(x_batch, training=True)   # forward pass inside tape
        loss = loss_fn(y_batch, predictions)           # scalar loss
    
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    return loss

Correctness Conditions of a Custom Training Step

Writing a loop is not enough — it must also be mathematically correct. A custom training step can run without errors and still be wrong. Gradients may be missing, losses improperly reduced, or variables silently excluded from updates.

What Makes a Training Step Valid?

Failure Modes Mapped to Execution Stages

flowchart TD
    A[Input] -->|Shape mismatch| B[Forward Pass]
    B -->|Detached tensors| C[Loss]
    C -->|NaNs / Improper reduction| D[Gradients]
    D -->|None gradients| E[Update]
    E -->|Silent no-op updates| F[Nothing learned]

• Input stage: Incorrect batch shapes propagate silently into later layers.
• Forward pass: Detached tensors or misplaced operations can remove parts of the graph from differentiation.
• Loss stage: Failing to aggregate properly results in incorrect gradient scaling.
• Gradient computation: None gradients signal variables were not involved in computation.
• Update step: Mismatched variable lists result in parameters that never change.

Understanding this mapping is what separates functional training from reliable training.

Demo: Train a Small Network End-to-End

The following notebook demonstrates training a small MLP using a fully custom loop on a synthetic dataset. The goal is to confirm that loss decreases, accuracy increases, and gradients remain finite — with an explicit debug workflow.

Setup

import os
import tensorflow as tf
import numpy as np

os.environ["TF_CPP_MIN_LOG_LEVEL"] = "1"
print(f"TensorFlow {tf.__version__}")

Data and Model

INPUT_DIM = 8
NUM_CLASSES = 3
BATCH_SIZE = 32
NUM_SAMPLES = 400

def make_data(num_samples, input_dim, num_classes, batch_size, seed=42):
    tf.random.set_seed(seed)
    x = tf.random.normal([num_samples, input_dim])
    raw = tf.reduce_sum(x[:, :2], axis=1) * 2
    y = tf.cast(raw, tf.int32) % num_classes
    ds = tf.data.Dataset.from_tensor_slices((x, y))
    return ds.shuffle(num_samples).batch(batch_size)

def build_model(input_dim, num_classes):
    return tf.keras.Sequential([
        tf.keras.layers.Dense(32, activation="relu", input_shape=(input_dim,)),
        tf.keras.layers.Dense(16, activation="relu"),
        tf.keras.layers.Dense(num_classes),
    ])

Minimal Debug Workflow

def check_shapes(x_batch, y_batch, predictions):
    assert x_batch.shape[1] == INPUT_DIM, "Input shape mismatch"
    assert predictions.shape[1] == NUM_CLASSES, "Output shape mismatch"
    print(f"  [shape check] x={x_batch.shape}, y={y_batch.shape}, pred={predictions.shape}")

def check_nans(loss, gradients):
    if tf.math.is_nan(loss):
        print("  [WARNING] NaN loss detected!")
    for g in gradients:
        if g is not None and tf.reduce_any(tf.math.is_nan(g)):
            print("  [WARNING] NaN gradient detected!")

LOSS_THRESHOLD = 5.0
def check_exploding(loss):
    if loss.numpy() > LOSS_THRESHOLD:
        print(f"  [WARNING] Exploding loss: {loss.numpy():.4f}")

Custom Training Loop

def train_one_epoch(model, optimizer, loss_fn, dataset, debug_first_step=True):
    loss_metric = tf.keras.metrics.Mean()
    acc_metric = tf.keras.metrics.SparseCategoricalAccuracy()

    for step, (x_batch, y_batch) in enumerate(dataset):
        with tf.GradientTape() as tape:
            predictions = model(x_batch, training=True)
            loss = loss_fn(y_batch, predictions)

        gradients = tape.gradient(loss, model.trainable_variables)
        optimizer.apply_gradients(zip(gradients, model.trainable_variables))

        if debug_first_step and step == 0:
            check_shapes(x_batch, y_batch, predictions)
            check_nans(loss, gradients)
            check_exploding(loss)

        loss_metric.update_state(loss)
        acc_metric.update_state(y_batch, predictions)

    return loss_metric.result().numpy(), acc_metric.result().numpy()

# Run training
dataset = make_data(NUM_SAMPLES, INPUT_DIM, NUM_CLASSES, BATCH_SIZE)
model = build_model(INPUT_DIM, NUM_CLASSES)
optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3)
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)

EPOCHS = 5
for epoch in range(1, EPOCHS + 1):
    loss, acc = train_one_epoch(model, optimizer, loss_fn, dataset,
                                debug_first_step=(epoch == 1))
    print(f"Epoch {epoch}/{EPOCHS} — loss: {loss:.4f}, acc: {acc:.4f}")

Expected Output

Epoch 1/5 — loss: 1.1200, acc: 0.4025
Epoch 2/5 — loss: 1.0900, acc: 0.4312
Epoch 3/5 — loss: 1.0700, acc: 0.4487
Epoch 4/5 — loss: 1.0500, acc: 0.4631
Epoch 5/5 — loss: 1.0200, acc: 0.4800

Loss decreasing and accuracy increasing confirms the loop is doing the right thing — gradients flowing, optimizer updating weights. No guesswork; the execution is fully visible.

Module 2: Optimization Controls — Clipping and Schedules

Why Training Becomes Unstable

Optimization controls directly affect training stability. The key question is: Is the update I am applying actually safe? Training can become unstable even when the architecture and data look fine. Most failures come from optimization dynamics — step size, gradient spikes, or sharp loss regions.

Sources of Gradient Instability

mindmap
  root((Instability))
    Deep Networks
      Multiplied Jacobians across layers
      Vanishing or exploding gradients
    Poor Initialization
      Saturated activations
      Gradient collapse
    Noisy Data
      High-variance batches
      Gradient jitter
    Large Learning Rate
      Overshoots repeatedly
      Early divergence
    Sharp Loss Regions
      High curvature amplifies differences
      Step outcomes unpredictable

1. Deep networks multiply Jacobians across many layers, so gradients can either shrink to near zero or blow up as depth increases. The same learning rate that works for a shallow model can collapse at a deeper one.

2. Poor initialization sets the starting scale of activations and gradients. If weights are too large, activations saturate and gradients become unstable. If they start too small, gradients vanish and learning stalls. The loop may still run, but updates become meaningless.

3. Noisy data — labels or high-variance batches — create gradient jitter even if the model is correct. The gradients bounce around because the signal-to-noise ratio is low. This often shows up as unstable loss curves that never settle.

4. Large learning rate controls how much you trust each gradient step. If step size is too aggressive, the optimizer overshoots repeatedly and the system can diverge fast, especially early when weights are uncalibrated.

5. Sharp loss regions — highly curved regions of the loss surface — amplify small numerical differences. The same gradient magnitude can produce very different outcomes step-by-step.

The key idea: instability usually has a traceable cause, and once you name it, you can choose the right control instead of guessing.

Learning Rate as a Stability Factor

Implementing Gradient Clipping and LR Schedules

Gradient Clipping: Norm vs. Value

Two practical controls directly shape update behavior and together reduce the probability that one bad step destroys the run.

graph TD
    A[Gradient Vector] --> B{Clipping Type?}
    B -->|By Norm| C[Scale entire gradient vector]
    B -->|By Value| D[Cap individual gradient elements]
    C --> E[Direction preserved, magnitude limited]
    D --> F[Direction may change, outliers capped]
    E --> G[Use for systemic / global instability]
    F --> H[Use for localized spikes]

def clip_gradients_by_norm(gradients, clip_norm=1.0):
    clipped, global_norm = tf.clip_by_global_norm(gradients, clip_norm)
    return clipped, global_norm

def clip_gradients_by_value(gradients, clip_value=0.5):
    return [tf.clip_by_value(g, -clip_value, clip_value) for g in gradients]

Learning Rate Schedules

Schedules shape update size over time, not just a single step. They are essentially controlled step-size policies designed to keep updates safe early and effective later.

Warmup + Cosine Decay

def warmup_cosine_lr(step, warmup_steps, total_steps, base_lr):
    if step < warmup_steps:
        # Linear warmup from 0 to base_lr
        return base_lr * (step / warmup_steps)
    else:
        # Cosine decay from base_lr to 0
        progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)
        return base_lr * 0.5 * (1.0 + tf.math.cos(tf.constant(3.14159) * progress))

Cyclic LR (Triangular)

def cyclic_lr(step, base_lr=1e-4, max_lr=1e-2, cycle_length=200):
    cycle = tf.math.floor(1 + step / (2 * cycle_length))
    x = tf.abs(step / cycle_length - 2 * cycle + 1)
    return base_lr + (max_lr - base_lr) * tf.maximum(0.0, 1.0 - x)

Schedules should be applied per step, not per epoch. Per-step scheduling keeps behavior consistent regardless of dataset size and batch size.

Applying Controls Inside the Loop

@tf.function
def train_step_with_controls(model, optimizer, loss_fn, x_batch, y_batch,
                              clip_norm=1.0, use_clipping=True):
    with tf.GradientTape() as tape:
        predictions = model(x_batch, training=True)
        loss = loss_fn(y_batch, predictions)          # must be scalar

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

    # Apply clipping BEFORE updates
    if use_clipping:
        gradients, global_norm = tf.clip_by_global_norm(gradients, clip_norm)
    else:
        global_norm = tf.linalg.global_norm(gradients)

    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    return loss, global_norm

The six controls that belong inside the loop:

1. Compute gradients explicitly (for visibility)
2. Apply clipping immediately after gradients are produced
3. Update learning rate per step to match the planned curve
4. Log gradient norms (your stability dashboard)
5. Log learning rate values per step (schedule verification)
6. Validate stability signals — gradients finite, loss finite, updates happening, norms within range

Tuning and Validating Stabilizers in Production

Choosing a Clipping Threshold

The objective is not to change the loss — it is to constrain the step size when it becomes unsafe.

Empirical approach: Log global gradient norms for a short run (hundreds to a few thousand steps). Pick a threshold that sits above the normal range but below extreme spikes.

Schedule and Clipping Interaction Patterns

Verification Checklist for Stable Optimization

flowchart TD
    A[Compute Loss - scalar] --> B[Compute Gradients]
    B --> C[Log Global Grad Norm]
    C --> D[Log Effective LR per step]
    D --> E[Apply Clipping if enabled]
    E --> F[Apply Optimizer Update]
    F --> G{Validate}
    G -->|Loss finite?| H[Continue]
    G -->|Grads finite?| H
    G -->|Updates non-zero?| H
    G -->|Norms in range?| H
    G -->|Any failure| I[Fail fast — stop training]

If validation fails, fail fast. Do not keep training and hope it recovers.

Demo: Stabilize Training with Clipping and Schedule

Four training configurations are compared on the same model and data, altering gradient and LR behavior:

def train_epoch(model, optimizer, loss_fn, dataset, step_counter,
                clip_mode=None, clip_value=1.0, lr_schedule_fn=None):
    loss_metric = tf.keras.metrics.Mean()
    acc_metric = tf.keras.metrics.SparseCategoricalAccuracy()
    grad_norms = []
    lr_values = []

    for x_batch, y_batch in dataset:
        step = step_counter[0]

        # Update LR from schedule
        if lr_schedule_fn is not None:
            new_lr = lr_schedule_fn(step)
            optimizer.learning_rate.assign(float(new_lr))
        lr_values.append(float(optimizer.learning_rate))

        with tf.GradientTape() as tape:
            preds = model(x_batch, training=True)
            loss = loss_fn(y_batch, preds)

        grads = tape.gradient(loss, model.trainable_variables)

        if clip_mode == "norm":
            grads, gn = tf.clip_by_global_norm(grads, clip_value)
        elif clip_mode == "value":
            grads = [tf.clip_by_value(g, -clip_value, clip_value) for g in grads]
            gn = tf.linalg.global_norm(grads)
        else:
            gn = tf.linalg.global_norm(grads)

        grad_norms.append(float(gn))
        optimizer.apply_gradients(zip(grads, model.trainable_variables))

        loss_metric.update_state(loss)
        acc_metric.update_state(y_batch, preds)
        step_counter[0] += 1

    return (loss_metric.result().numpy(), acc_metric.result().numpy(),
            grad_norms, lr_values)

Comparison Results

Overall findings:

• Gradient clipping by norm ensures smoother, more stable training.
• Value clipping works but is less consistent.
• No clipping can yield competitive short-term results but risks instability over longer runs.
• Cyclic LR adds exploration but may need clipping to stabilize gradient swings.

Module 3: Advanced Loop Techniques — Mixed Precision and Nested Tape

Mixed Precision Mechanics That Matter

Mixed precision and nested GradientTape are two tools that unlock performance and expressiveness, but both can break training if the mechanics are not handled carefully. The goal is to keep speed gains while keeping gradients correct and stable.

What Mixed Precision Changes in Practice

graph TD
    subgraph "Compute Path - FP16 for speed"
        A[Forward Pass] --> B[MatMul / Conv ops]
        B --> C[FP16 Activations]
    end
    subgraph "Stability Path - FP32 for correctness"
        D[Loss Value] --> E[Loss Scaling]
        E --> F[Gradient Reduction]
    end
    C --> D

Speed-up depends heavily on hardware. GPUs with Tensor Cores (Nvidia RTX, A100) and TPUs have dedicated pathways that make FP16/BF16 compute substantially faster than FP32. Without that support, the gain can be small or inconsistent.

Loss Scaling and the Underflow Problem

FP16 has a much smaller dynamic range than FP32. If gradients become very small (common in deep models or late-stage training), those values can underflow to zero in FP16.

flowchart LR
    A[Small Gradients] --> B{FP16 Range?}
    B -->|Too small| C[Underflow to 0]
    C --> D[No parameter updates]
    D --> E[Training stalled silently]
    B -->|In range| F[Normal gradient flow]

The fix — loss scaling:

1. Multiply the loss by a large scale factor before computing gradients. This inflates gradients into the representable FP16 range.
2. After gradients are computed, unscale them back to their correct magnitude before applying updates.
3. NaN detection: if scaling is too aggressive, intermediate values can blow up to inf or NaN. Practical mixed precision uses NaN checks and adjusts the scale dynamically.

Loss scaling is not optional in FP16 pipelines — it is the stability mechanism that keeps gradients numerically representable.

Nested GradientTape and Higher-Order Differentiation

Beyond first-order gradients, some training objectives require differentiating through the gradient computation itself. This is what nested GradientTape enables.

How Nested Tapes Work

flowchart TD
    A[Inputs / Interpolated Samples] --> B[Inner GradientTape]
    B --> C[Forward Pass]
    C --> D[First-order Gradients w.r.t. inputs]
    D --> E[Outer GradientTape]
    E --> F[Second-order Derivatives]
    F --> G[Gradient Penalty Term]
    G --> H[Total Loss]

• Inner tape records the forward computation so gradients can be taken with respect to model parameters or inputs, depending on the penalty being built.
• Outer tape records the process of forming those gradients. The gradient computation itself becomes differentiable.
• Second derivatives become possible because the outer tape watches how gradients are formed.
• Regularization terms defined in terms of gradients (e.g., “gradients should behave like x”) require differentiating through the gradient computation.

Nested tapes convert gradients from a byproduct into a first-class object that can be optimized.

Where Gradient Penalties Are Used

Gradient penalties enforce smoothness by discouraging the model from changing too sharply with respect to inputs. They reduce sensitivity and improve stability.

Standard high-level training flows cannot express these patterns cleanly — direct control of the steps is required so the penalty is computed correctly.

Engineering Pattern: Safe Mixed Precision and Nested Tape

Loss Scaling Implementation Pattern

from tensorflow.keras import mixed_precision

# Set global dtype policy
mixed_precision.set_global_policy("mixed_float16")

# Define base optimizer and wrap with loss scale optimizer
base_opt = tf.keras.optimizers.Adam(learning_rate=1e-3)
opt = mixed_precision.LossScaleOptimizer(base_opt)

@tf.function
def train_step_mixed(model, opt, loss_fn, x, y):
    with tf.GradientTape() as tape:
        y_pred = model(x, training=True)          # forward in FP16 kernels
        loss = loss_fn(y, y_pred)                 # scalar loss
        scaled_loss = opt.get_scaled_loss(loss)   # inflate for FP16 range

    scaled_grads = tape.gradient(scaled_loss, model.trainable_variables)
    grads = opt.get_unscaled_gradients(scaled_grads)  # restore true magnitude

    opt.apply_gradients(zip(grads, model.trainable_variables))
    return loss

The stable pattern:

1. Scale before gradients — prevent FP16 underflow turning small gradients into zeros.
2. Unscale before update — so the optimizer applies the true gradient magnitude.
3. Keep the step minimal and verifiable: loss finite, grads not NaN, norm sane, weights actually change.

Gradient Penalty via Nested Tape

def gradient_penalty(model, real_inputs, fake_inputs):
    # Interpolate between real and fake samples
    alpha = tf.random.uniform([tf.shape(real_inputs)[0], 1], 0.0, 1.0)
    interpolated = alpha * real_inputs + (1 - alpha) * fake_inputs

    with tf.GradientTape() as inner_tape:
        inner_tape.watch(interpolated)
        predictions = model(interpolated, training=True)
        scalar_output = tf.reduce_sum(predictions)  # collapse to scalar

    # First-order gradients with respect to the interpolated inputs
    grads = inner_tape.gradient(scalar_output, interpolated)

    # Squared L2 norm of the gradient, averaged across the batch
    grad_norm_sq = tf.reduce_sum(tf.square(grads), axis=1)
    penalty = tf.reduce_mean(grad_norm_sq)
    return penalty

Then in the outer training step:

@tf.function
def train_step_with_penalty(model, optimizer, loss_fn, x_real, y, x_fake,
                             penalty_weight=0.1):
    with tf.GradientTape() as outer_tape:
        predictions = model(x_real, training=True)
        main_loss = loss_fn(y, predictions)
        gp = gradient_penalty(model, x_real, x_fake)
        total_loss = main_loss + penalty_weight * gp   # combined objective

    grads = outer_tape.gradient(total_loss, model.trainable_variables)
    optimizer.apply_gradients(zip(grads, model.trainable_variables))
    return main_loss, gp

Nested Tape Correctness Checklist

Demo: Mixed Precision with Gradient Penalty

The demo combines mixed precision training with a gradient penalty using nested GradientTapes. Key steps:

# 1. Set mixed precision policy
try:
    mixed_precision.set_global_policy("mixed_float16")
    print(f"Active policy: {mixed_precision.global_policy().name}")
except Exception as e:
    print(f"Falling back to float32: {e}")

# 2. Build model and dataset
def make_data(num_samples=640, input_dim=8, num_classes=3, batch_size=32):
    tf.random.set_seed(42)
    x = tf.random.normal([num_samples, input_dim])
    y = tf.cast(tf.reduce_sum(x[:, :2], axis=1) * 2, tf.int32) % num_classes
    return tf.data.Dataset.from_tensor_slices((x, y)).shuffle(num_samples).batch(batch_size)

def build_model(input_dim=8, num_classes=3):
    return tf.keras.Sequential([
        tf.keras.layers.Dense(24, activation="relu", input_shape=(input_dim,)),
        tf.keras.layers.Dense(num_classes),  # raw logits
    ])

# 3. Training loop with numerics check
EPOCHS = 3
dataset = make_data()
model = build_model()
optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3)
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)

for epoch in range(1, EPOCHS + 1):
    epoch_losses, epoch_accs, gp_values = [], [], []
    for x_real, y_batch in dataset:
        x_fake = tf.random.normal(tf.shape(x_real))
        main_loss, gp = train_step_with_penalty(model, optimizer, loss_fn,
                                                 x_real, y_batch, x_fake)
        epoch_losses.append(float(main_loss))
        gp_values.append(float(gp))

    avg_loss = np.mean(epoch_losses)
    finite_ok = np.isfinite(avg_loss)
    print(f"Epoch {epoch}/{EPOCHS} | loss={avg_loss:.4f} | "
          f"GP_mean={np.mean(gp_values):.4f} | finite={finite_ok}")

print("numerics OK (no NaNs); gradient penalty computed. Done.")

Key observations:

• Mixed precision policy pushes matmul/conv into faster kernels.
• On CPU-only machines, FP16 policy is set but computation may not materially accelerate — it activates fully on GPU with Tensor Cores.
• The gradient penalty being non-NaN confirms it is being computed correctly through the nested tape structure.

Module 4: Custom Layers and Models — Subclassing for Real Architectures

When to Subclass Layer and Model

Custom architectures fail not because TensorFlow is missing features, but when execution contracts are violated. The focus shifts from assembling layers to controlling how state is created, tracked, and replayed, so custom components behave correctly during training and after saving.

The decision to subclass is triggered by one thing: you need behavior that the graph cannot express clearly with standard wiring.

Functional API vs. Subclassing

graph LR
    A[Functional API] -->|Static graph| B[Fixed data flow]
    A -->|Easy visualization| C[Inspectable structure]
    A -->|Best for| D[Standard models]
    E[Subclassing] -->|Dynamic execution| F[Runtime conditional logic]
    E -->|Full control| G[Arbitrary computation]
    E -->|Best for| H[Custom behavior]

The choice comes down to whether structure or flexibility is the dominant requirement.

Layer and Model Lifecycle Contracts

Four lifecycle stages that must be respected:

flowchart LR
    A[build - shape known, create variables] --> B[call - compute forward pass]
    B --> C[variables - tracked before execution]
    C --> D[learning step - optimizer updates correctly]

• build(input_shape) — where shapes become known and variables are created. Getting this right is critical because it establishes the parameters training will later update.
• call(inputs, training) — where computation actually happens. Keeping call free of side effects makes behavior predictable.
• variables — all trainable state must exist before execution reaches call. Variables created in the wrong place cause tracking to break silently.
• Learning step — if earlier contracts are respected, optimization behaves as expected. If not, large updates amplify hidden mistakes.

Implementing Custom Components

Custom Layer Structure

class SimpleNormLayer(tf.keras.layers.Layer):
    def __init__(self, eps=1e-5, **kwargs):
        super().__init__(**kwargs)
        self.eps = eps

    def build(self, input_shape):
        last_dim = input_shape[-1]
        # Add trainable weights using add_weight (correct tracking)
        self.gamma = self.add_weight(
            name="gamma", shape=(last_dim,),
            initializer="ones", trainable=True
        )
        self.beta = self.add_weight(
            name="beta", shape=(last_dim,),
            initializer="zeros", trainable=True
        )
        super().build(input_shape)

    def call(self, inputs, training=None):
        mean = tf.reduce_mean(inputs, axis=-1, keepdims=True)
        variance = tf.math.reduce_variance(inputs, axis=-1, keepdims=True)
        normalized = (inputs - mean) / tf.sqrt(variance + self.eps)
        return self.gamma * normalized + self.beta  # learned affine transform

Key points:

• build() is called automatically the first time the layer sees real data.
• gamma and beta are added via self.add_weight() — the correct mechanism for registering trainable variables.
• eps prevents division by zero when variance is extremely small.
• super().build(input_shape) marks the layer as built internally.

Custom Model Forward Logic

class CustomSmallModel(tf.keras.Model):
    def __init__(self, num_classes, **kwargs):
        super().__init__(**kwargs)
        self.norm_layer = SimpleNormLayer()          # custom layer as attribute
        self.dense1 = tf.keras.layers.Dense(24, activation="relu")
        self.dense2 = tf.keras.layers.Dense(num_classes)  # raw logits

    def call(self, inputs, training=None):
        x = self.dense1(inputs)
        if training:                                 # conditional branching
            x = self.norm_layer(x, training=training)
        return self.dense2(x)

Key points:

• Subclassing tf.keras.Model gives you a full model object that Keras can track, save, and integrate with optimizers.
• All forward computation lives in call().
• Conditional logic based on training flag — normalization is active only during training.
• Sublayers are registered as attributes (self.norm_layer, self.dense1, self.dense2) so Keras discovers their variables.

Production Readiness: Tracking and Saving Subclassed Components

Weight Tracking and State Management Rules

The critical warning: A model that runs can still silently skip updates. If tracking is broken, the tape can return None gradients, or the optimizer can receive an empty variable list — and the step still executes without raising an error.

Subclassed Component Lifecycle for Save/Load

flowchart TD
    A["__init__ stores configuration\nnum_units, dropout_rate, etc."] --> B["build(input_shape)\ncreates weights deterministically"]
    B --> C["call(inputs, training)\ndefines forward behavior consistently"]
    C --> D["get_config()\nserializes hyperparameters"]
    D --> E["save() / load()\npreserves config + weights + forward path"]
    E --> F["Reload verification\nsame input → same output"]

Practical Serialization Checklist

class CustomDenseLayer(tf.keras.layers.Layer):
    def __init__(self, units, activation="relu", **kwargs):
        super().__init__(**kwargs)
        self.units = units                           # store hyperparameters
        self.activation = activation

    def build(self, input_shape):
        self.kernel = self.add_weight(
            "kernel", shape=(input_shape[-1], self.units), trainable=True
        )
        self.bias = self.add_weight(
            "bias", shape=(self.units,), initializer="zeros", trainable=True
        )
        super().build(input_shape)

    def call(self, inputs, training=None):
        out = tf.matmul(inputs, self.kernel) + self.bias
        if self.activation == "relu":
            return tf.nn.relu(out)
        return out

    def get_config(self):                            # enable serialization
        config = super().get_config()
        config.update({"units": self.units, "activation": self.activation})
        return config

# Verify serialization
layer = CustomDenseLayer(32)
dummy = tf.ones([1, 16])
layer(dummy)                                        # trigger build

config = layer.get_config()
reconstructed = CustomDenseLayer.from_config(config)
reconstructed(dummy)

# Outputs should match within numeric tolerance
original_out = layer(dummy)
reconstructed_out = reconstructed(dummy)
# NOTE: weights differ after fresh initialization; test after checkpoint save/load

Serialization truth test: After a reload, the same input should produce the same output (within expected numeric tolerance). If outputs shift, either config was not preserved, weights did not map correctly, or the forward path depends on state that is not being tracked consistently.

Demo: Plug Custom Components into Training

Training Loop with Custom Layer Verification

# Build dataset
dataset = make_data(num_samples=800, input_dim=8, num_classes=3, batch_size=32)

# Instantiate model
model = CustomSmallModel(num_classes=3)
optimizer = tf.keras.optimizers.Adam(learning_rate=1e-2)
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)

# Force build by running a dummy forward pass
dummy_input = tf.zeros([1, 8])
_ = model(dummy_input, training=True)

# Capture pre-training weights
gamma_before = model.norm_layer.gamma.numpy().copy()
beta_before = model.norm_layer.beta.numpy().copy()

# Training loop
EPOCHS = 4
for epoch in range(1, EPOCHS + 1):
    loss_metric = tf.keras.metrics.Mean()
    acc_metric = tf.keras.metrics.SparseCategoricalAccuracy()

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

        gradients = tape.gradient(loss, model.trainable_variables)
        optimizer.apply_gradients(zip(gradients, model.trainable_variables))
        loss_metric.update_state(loss)
        acc_metric.update_state(y_batch, predictions)

    print(f"Epoch {epoch}/{EPOCHS} — "
          f"loss: {loss_metric.result():.4f}, "
          f"acc: {acc_metric.result():.4f}")

# Verify custom layer weights actually changed
gamma_after = model.norm_layer.gamma.numpy()
beta_after = model.norm_layer.beta.numpy()
gamma_changed = not np.allclose(gamma_before, gamma_after)
beta_changed = not np.allclose(beta_before, beta_after)
print(f"gamma_changed: {gamma_changed} | beta_changed: {beta_changed}")

What this verifies:

• gamma_changed: True and beta_changed: True proves the custom normalization layer actively participated in the learning process.
• Output shape (2, 3) for a 2-sample batch confirms the model’s last dimension is correctly aligned with num_classes.
• Loss decreasing and accuracy increasing over epochs confirms the custom loop drives correct learning.

Module 5: Custom Training Components — Losses, Metrics, Optimizers, and train_step()

Designing Custom Losses and Metrics

Training becomes reliable when three pieces are treated as first-class: losses, metrics, and optimizers. When any one of them is hand-waved, you can still get runs that look normal but behave wrong.

In production, accuracy plus cross-entropy is only a narrow slice of what systems optimize. You might care about calibration, constraint violations, latency-aware penalties, fairness terms, reconstruction fidelity, or stability of intermediate representations.

When you move beyond default objectives:

• A custom metric becomes your instrument panel.
• A custom loss becomes your control target.
• If either is poorly defined — wrong reduction, wrong dtype, wrong weighting — the training loop can look healthy while silently optimizing the wrong thing.

Multi-term Loss Composition

graph TD
    A[Primary Loss] --> E[Total Objective - scalar]
    B[Regularization Term] --> E
    C[Auxiliary Penalty] --> E
    D[Constraint Penalties] --> E
    E --> F[Gradient Computation]

The technical rule: the total must reduce to a single scalar per step with a reduction you can justify. If aggregation is inconsistent, gradient magnitude will scale incorrectly with batch size.

Stateful vs. Stateless Metrics

Practical rule: Use stateless for fast per-batch signals. Use stateful when you need trustworthy trends. Treat state updates and resets as part of training logic — not an afterthought.

Custom Optimizers and train_step()

Overriding Training Logic

Overriding train_step() lets you add specialized behavior while keeping model.fit() as the orchestration layer:

flowchart TD
    A[Forward Pass] --> B[Loss Compute]
    B --> C[Gradient Update]
    C --> D[Metric Update]
    D --> E[Return metrics dict]

Each stage can be customized:

• Forward pass — can expose intermediate activations or use multiple forward passes.
• Loss compute — architecture-driven multi-term objectives.
• Gradient update — clipping, multiple optimizers, conditional updates.
• Metric update — timing matters; update after masking and sample weighting.

When you override training logic, you are not just customizing code — you are defining the optimization algorithm.

Custom Optimizer Structure

class CustomMomentumOptimizer(tf.keras.optimizers.Optimizer):
    def __init__(self, learning_rate=0.01, momentum=0.9, **kwargs):
        super().__init__(name="CustomMomentum", **kwargs)
        self._learning_rate = self._build_learning_rate(learning_rate)
        self.momentum = momentum

    def build(self, var_list):
        """Initialize optimizer state variables once per trainable variable."""
        super().build(var_list)
        self._velocity = [
            self.add_variable_from_reference(var, "velocity")
            for var in var_list
        ]

    def update_step(self, gradient, variable, learning_rate):
        """Core update rule — called per variable per step."""
        idx = self._get_variable_index(variable)
        v = self._velocity[idx]
        # Momentum update rule: v = momentum * v - lr * gradient
        v.assign(self.momentum * v - learning_rate * gradient)
        variable.assign_add(v)

    def get_config(self):
        config = super().get_config()
        config.update({
            "learning_rate": self._serialize_hyperparameter("learning_rate"),
            "momentum": self.momentum,
        })
        return config

Custom optimizer requirements:

1. Subclass tf.keras.optimizers.Optimizer — keeps variables tracked, checkpoints working, training loops compatible.
2. Define the update rule — the core behavior: how a gradient becomes a parameter change.
3. Manage optimizer state — momentum buffers, second-moment accumulators, EMAs must be created once per variable and updated deterministically.
4. Keep it compatible with GradientTape outputs — preserve the one-to-one mapping between gradients and variables.
5. Preserve ecosystem behavior — checkpointing, serialization, mixed precision, distributed strategies.

Validating Custom Components Before You Trust Them

Component Validation Flow

flowchart TD
    A[Unit-level Checks] --> B[Numerical Sanity Checks]
    B --> C[Integration in train_step]
    C --> D[Regression after Refactors]
    D --> E{Key Outcomes}
    E --> F[Loss scale stable]
    E --> G[Metrics reset correctly]
    E --> H[Optimizer updates change variables]

Unit-level checks:

• Loss reduces to one scalar
• Tensor shapes align with model expectations
• Metrics have state variables when they should
• Optimizer can consume the exact grad/var pairs you plan to apply

Numerical sanity checks:

• Loss stays finite
• Gradients are finite
• Gradient norms sit in a reasonable band for your model scale
• If mixed precision is enabled, dtype behavior is stable

Integration in train_step:

• Forward pass, loss assembly, metric updates, optimizer application must happen in a consistent sequence
• The same logic that worked in isolation still behaves correctly when everything is wired together

Loss and Metric Sanity Checks

Mean vs. Sum reduction matters: Mean keeps gradient scale stable as batch size changes. Sum makes updates grow with batch size, which can create hardware-dependent instability.

Optimizer Verification Loop

def verify_optimizer(optimizer_fn, variable_init=5.0):
    """Verify optimizer correctness on a toy loss."""
    var = tf.Variable(variable_init)

    # Toy loss: minimize x^2, minimum at x=0
    # Gradient at x: 2x
    with tf.GradientTape() as tape:
        loss = var ** 2

    grad = tape.gradient(loss, var)

    # Verify gradient exists and has correct sign
    assert grad is not None, "Gradient is None — variable not in tape scope"
    assert grad.numpy() > 0, "Gradient should be positive for x > 0"
    assert np.isfinite(grad.numpy()), "Gradient is not finite"

    var_before = var.numpy()
    optimizer_fn([grad], [var])
    var_after = var.numpy()

    # Verify variable moved in the correct direction
    assert var_after < var_before, "Variable should decrease (gradient descent)"
    print(f"  var: {var_before:.4f} -> {var_after:.4f} | "
          f"grad: {grad.numpy():.4f} | loss: {loss.numpy():.4f}")
    print("  Optimizer verification: PASSED")

Optimizer Correctness Checklist

Reload check: After reloading, run a few steps on the same batch and confirm behavior matches pre-save within tolerance. If it does not match, momentum or accumulators were lost.

Demo: Specialized Training with train_step()

CustomTrainStepModel

class CustomTrainStepModel(tf.keras.Model):
    def __init__(self, num_classes, **kwargs):
        super().__init__(**kwargs)
        self.d1 = tf.keras.layers.Dense(24, activation="relu")
        self.d2 = tf.keras.layers.Dense(num_classes)  # raw logits

        # Named metrics — appear in model.fit() logs automatically
        self.loss_tracker = tf.keras.metrics.Mean(name="loss")
        self.acc_metric = tf.keras.metrics.SparseCategoricalAccuracy(name="accuracy")
        self.aux_tracker = tf.keras.metrics.Mean(name="aux_penalty")

    def call(self, inputs, training=None):
        x = self.d1(inputs)
        return self.d2(x)

    def train_step(self, data):
        x, y = data

        with tf.GradientTape() as tape:
            logits = self(x, training=True)

            # Main classification loss
            main_loss = tf.reduce_mean(
                tf.keras.losses.sparse_categorical_crossentropy(
                    y, logits, from_logits=True
                )
            )
            # Auxiliary L2 penalty on logit magnitude
            aux_penalty = 0.01 * tf.reduce_mean(tf.square(logits))

            total_loss = main_loss + aux_penalty

        grads = tape.gradient(total_loss, self.trainable_variables)
        self.optimizer.apply_gradients(zip(grads, self.trainable_variables))

        # Update metrics (main loss tracked separately from total)
        self.loss_tracker.update_state(main_loss)
        self.acc_metric.update_state(y, logits)
        self.aux_tracker.update_state(aux_penalty)

        return {
            "loss": self.loss_tracker.result(),
            "accuracy": self.acc_metric.result(),
            "aux_penalty": self.aux_tracker.result(),
        }

    @property
    def metrics(self):
        # Tell Keras which metrics to reset at each epoch start
        return [self.loss_tracker, self.acc_metric, self.aux_tracker]

Running with model.fit()

# Synthetic dataset
x = tf.random.normal([640, 8], seed=42)
y = tf.cast(tf.reduce_sum(x[:, :2], axis=1) * 2, tf.int32) % 3
dataset = tf.data.Dataset.from_tensor_slices((x, y)).shuffle(320).batch(32)

# Compile without loss argument — train_step handles it internally
model = CustomTrainStepModel(num_classes=3)
model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=1e-2))

history = model.fit(dataset, epochs=6)

Output includes per-epoch: loss, accuracy, aux_penalty — all from our custom metrics.

Key design advantages:

• Keras still handles epochs, batches, and progress bars.
• Custom loss components influence optimization.
• Custom metrics appear in history and progress output automatically.
• Clean separation: architecture in call(), training logic in train_step(), metric tracking in metrics property.

This pattern is the right abstraction level when you need extra loss terms, per-batch custom logic, or custom metrics — without writing an entirely manual training loop from scratch.

Module 6: Distributed Training — Strategies That Ship

Strategy Selection and Trade-offs

At this point, training moves beyond a single device into environments where coordination matters as much as computation. Every strategy reflects a balance between simplicity, scalability, and operational cost.

Distributed Strategy Comparison

graph TD
    subgraph "MirroredStrategy"
        A1[Single Machine]
        A2[Multiple GPUs]
        A3[Synchronous replicas]
        A4[Shared data access]
    end
    subgraph "MultiWorkerMirroredStrategy"
        B1[Multiple Machines]
        B2[Synchronous workers]
        B3[Requires TF_CONFIG]
        B4[Scales horizontally]
    end
    subgraph "TPUStrategy"
        C1[TPU Pods]
        C2[Hardware-parallel cores]
        C3[TPU-specific setup]
        C4[Maximum throughput]
    end

Gradient Aggregation Models

How gradients are combined across devices determines correctness and communication cost:

TPU-Specific Checklist

• Use TPUStrategy with a valid TPU resolver; create model and dataset inside strategy.scope()
• Prefer fixed batch sizes and avoid dynamic shapes (so the step compiles efficiently)
• Wrap the training step in tf.function for graph execution on TPU
• Use strategy.run() for the step and scale loss by global batch size so gradients aggregate correctly

Correctness and Performance in Distributed Training

Global Batch Size and Gradient Scaling

flowchart LR
    A["Global Batch = Batch per replica × Number of replicas"] --> B{LR Adjustment?}
    B -->|Scale LR with global batch| C["Linear scaling\nBigger batch → more stable gradient\n→ larger step affordable"]
    B -->|Keep LR fixed| D["Conservative\nUse warmup + clipping\nfor stability"]
    B -->|No adjustment| E["Silent failure\nUpdate magnitude changes\nstability breaks"]

The key risk: If global batch size changes without adjusting learning rate behavior, update magnitude changes and stability can break. This might not crash immediately — instead you see slower learning, oscillation, or divergence after thousands of steps.

Treat global batch size like a change to the optimizer, not just the hardware. If you do not tune for it, you are training a different system even though the code looks the same.

Distributed Step Semantics

flowchart TD
    A[Replica 0: Forward + Loss] --> C[Cross-Replica Aggregation sum or mean]
    B[Replica 1: Forward + Loss] --> C
    C --> D[Single Optimizer Update - global model]
    D --> E[All replicas: same updated weights]

The key point: Even though computation happens per replica, there is one conceptual model update per step. Correctness depends on treating the optimizer update as global and ensuring the aggregated gradient matches the intended scaling.

Sum vs. Mean during aggregation:

• Mean tends to preserve “one step equals one step” behavior across replica counts.
• Sum often requires explicit scaling decisions and effectively changes the learning rate unless compensated.

Performance Bottleneck Checklist

Converting Code to Distributed Execution

Strategy Scope Execution

flowchart TD
    A[Enter strategy.scope] --> B[Create model variables - replicated]
    B --> C[Create optimizer variables - replicated]
    C --> D[Distribute dataset - sharded per replica]
    D --> E[strategy.run: per-replica forward + backward]
    E --> F[strategy.reduce: aggregate losses / metrics]
    F --> G[Single optimizer update - global]

The key sequence when converting a working custom loop to distributed:

1. Enter strategy.scope() — variable creation happens inside here.
2. Variables are created inside the scope — correctly replicated or partitioned across devices.
3. Computation runs in the replica context — each device performs forward/backward independently.
4. Cross-replica synchronization — gradients or results are combined so global model state stays consistent.

Distributed Loop Skeleton

import tensorflow as tf

# Choose strategy — falls back to CPU if no GPU
try:
    strategy = tf.distribute.MirroredStrategy()
except:
    strategy = tf.distribute.OneDeviceStrategy("/cpu:0")

print(f"Number of devices: {strategy.num_replicas_in_sync}")

# Create model and optimizer INSIDE strategy.scope()
with strategy.scope():
    model = tf.keras.Sequential([
        tf.keras.layers.Input(shape=(8,)),
        tf.keras.layers.Dense(24, activation="relu"),
        tf.keras.layers.Dense(3),
    ])
    optimizer = tf.keras.optimizers.Adam(learning_rate=1e-2)
    loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(
        from_logits=True, reduction=tf.keras.losses.Reduction.NONE
    )

# Create regular dataset, then distribute it
dataset = make_data(num_samples=640, input_dim=8, num_classes=3, batch_size=32)
dist_ds = strategy.experimental_distribute_dataset(dataset)

@tf.function
def train_step(inputs):
    x, y = inputs

    with tf.GradientTape() as tape:
        predictions = model(x, training=True)
        # Per-example loss (no reduction yet)
        per_example_loss = loss_fn(y, predictions)
        # Scale correctly for global batch
        loss = tf.nn.compute_average_loss(
            per_example_loss,
            global_batch_size=32
        )

    gradients = tape.gradient(loss, model.trainable_variables)
    # Under MirroredStrategy: gradients are auto-aggregated before update
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    return loss

def run_epoch(dist_dataset):
    total_loss = 0.0
    num_batches = 0

    for dist_batch in dist_dataset:
        per_replica_losses = strategy.run(train_step, args=(dist_batch,))
        # Aggregate per-replica losses into one scalar
        batch_loss = strategy.reduce(
            tf.distribute.ReduceOp.MEAN,
            per_replica_losses,
            axis=None
        )
        total_loss += batch_loss.numpy()
        num_batches += 1

    return total_loss / num_batches

# Training
EPOCHS = 3
for epoch in range(1, EPOCHS + 1):
    avg_loss = run_epoch(dist_ds)
    print(f"Epoch {epoch}/{EPOCHS} | avg_loss: {avg_loss:.4f}")

Multi-Worker Setup

For multiple machines, use MultiWorkerMirroredStrategy. The key extra piece is TF_CONFIG:

import os, json

# Set on each machine before importing TensorFlow
os.environ["TF_CONFIG"] = json.dumps({
    "cluster": {
        "worker": ["worker0.example.com:12345", "worker1.example.com:23456"]
    },
    "task": {"type": "worker", "index": 0}  # index=1 on the second machine
})

strategy = tf.distribute.MultiWorkerMirroredStrategy()
# Rest of the code is identical to MirroredStrategy

The same pattern of strategy.scope(), strategy.experimental_distribute_dataset(), and strategy.run() applies. The only difference is that replicas now span multiple workers.

Demo: Distributed Custom Loop

The demo demonstrates the full distributed custom loop, confirming:

• Losses decrease across epochs under distributed execution.
• strategy.run() dispatches to each replica in parallel.
• strategy.reduce(ReduceOp.MEAN, ...) produces a single scalar per batch.
• On CPU-only machines, MirroredStrategy falls back gracefully with one replica.

Expected output:

Number of devices: 1
Epoch 1/3 | avg_loss: 1.0872
Epoch 2/3 | avg_loss: 1.0614
Epoch 3/3 | avg_loss: 1.0389
Distributed code path ran successfully. Synchronized updates across replicas.
Done.

The three additions that enable distribution:

1. A tf.distribute.Strategy instance and strategy.scope()
2. A distributed dataset via strategy.experimental_distribute_dataset()
3. strategy.run() plus strategy.reduce() around your step function

With those in place, a custom loop can scale from single CPU to multiple GPUs on one machine to multiple workers and even to TPUs, all following the same pattern.

Module 7: Architecture Spotlights — GNNs, VAEs, GANs, RL, and scikit-learn Integration

Pattern Map for Advanced Workflows

Here we move beyond conventional architectures and examine models that fundamentally reshape how training logic is structured: graph neural networks, variational autoencoders, generative adversarial networks, and reinforcement learning systems.

What Changes in Training Loops for Advanced Architectures

mindmap
  root((Advanced Training)
    Data Structure
      Fixed tensors → Graphs or env states
      Latent distributions
    Loss Design
      Single objective → Composite/probabilistic
      KL divergence + reconstruction
      Discriminator + generator
    Update Order
      Uniform → Alternating/conditional
      Phase A then Phase B
    Optimizer Flow
      Automatic → Manual scheduling
      Staged updates per architecture

Architecture determines execution structure, not the other way around.

Architecture-Specific Constraints

Architecture-Driven Training Step Patterns

Multi-Objective Loss Structure

graph TD
    A[Primary Objective] --> E[Total Loss - scalar]
    B[Regularization Terms] --> E
    C[Auxiliary Objectives] --> E
    D[Constraint Penalties] --> E
    E --> F[Single Backward Pass]

• Primary objective — the core signal: classification error, reconstruction error, value error.
• Regularization terms — prevent trivial or unstable solutions: weight decay, norm constraints, spectral penalties, gradient penalties.
• Auxiliary objectives — improve learning dynamics: feature matching, contrastive losses, self-supervised heads.
• Constraint penalties — enforce architectural rules: KL constraints, trust region penalties, consistency constraints.

Once loss is composed this way, stability depends on consistent scaling and reduction every step.

Multi-Phase Updates Inside One Step

flowchart TD
    A[Phase A: Update subset of variables\ne.g., discriminator or critic] --> B[Phase B: Update different subset\ne.g., generator, conditioned on Phase A]
    B --> C[Optional Phase C: Update target networks or EMA weights]
    C --> D[Next training step]

Key point: execution order is part of correctness. With staged updates, the same computations in a different order can mean a different training algorithm. This is a correctness constraint, not a style choice.

Where multi-phase updates appear:

Building GNN Layers and Models

GNN Message Passing Flow

flowchart LR
    A[Node Features] --> B[Edge Messages\ncomputed from connected nodes]
    B --> C[Aggregation\nsum, mean, or attention]
    C --> D[Node Update\ntransform aggregated signal]
    D -->|Next layer| A

This sequence repeats across layers, forming iterative relational reasoning. Each layer allows nodes to gather information from farther neighbors.

Why GNNs Require Custom Loops

A minimal GNN message-passing layer:

class GNNMessagePassingLayer(tf.keras.layers.Layer):
    def __init__(self, units, **kwargs):
        super().__init__(**kwargs)
        self.units = units
        self.message_fn = tf.keras.layers.Dense(units, activation="relu")
        self.update_fn = tf.keras.layers.Dense(units, activation="relu")

    def build(self, input_shape):
        super().build(input_shape)

    def call(self, node_features, adjacency_matrix, training=None):
        # adjacency_matrix: [num_nodes, num_nodes] sparse or dense
        # Compute messages: for each edge, transform source node features
        messages = self.message_fn(node_features)  # [N, units]

        # Aggregate: sum messages from neighbors
        # adjacency_matrix[i, j] = 1 if edge j->i exists
        aggregated = tf.matmul(adjacency_matrix, messages)  # [N, units]

        # Update: combine with current node features
        updated = self.update_fn(
            tf.concat([node_features, aggregated], axis=-1)
        )
        return updated

A full GNN model with multiple message passing layers:

class SimpleGNN(tf.keras.Model):
    def __init__(self, num_layers, hidden_units, num_classes, **kwargs):
        super().__init__(**kwargs)
        self.mp_layers = [
            GNNMessagePassingLayer(hidden_units)
            for _ in range(num_layers)
        ]
        self.readout = tf.keras.layers.Dense(num_classes)

    def call(self, node_features, adjacency_matrix, training=None):
        x = node_features
        for layer in self.mp_layers:
            x = layer(x, adjacency_matrix, training=training)
        # Graph-level readout: mean over nodes
        graph_embedding = tf.reduce_mean(x, axis=0, keepdims=True)
        return self.readout(graph_embedding)

VAEs and GANs with Custom Loops

Dual-Loop and ELBO Training

VAEs and GANs both require multiple objectives and explicit control over update timing.

graph TD
    subgraph VAE
        V1[Reconstruction Loss] --> V3[ELBO = recon - KL]
        V2[KL Divergence] --> V3
        V3 --> V4[Single backward pass through encoder + decoder]
    end
    subgraph GAN
        G1[Discriminator Loss] --> G2[Update D n times]
        G3[Generator Loss] --> G4[Update G once]
        G2 --> G4
    end

Why model.fit() fails for these architectures:

• Multiple losses per step cannot always be reduced into a single scalar objective.
• Alternating optimization phases break the assumption of uniform updates.
• Conditional update logic introduces branching inside the training loop.
• Non-scalar intermediate objectives prevent automatic gradient handling.
• Tight control over execution order requires manual orchestration.

VAE Training Step

@tf.function
def vae_train_step(encoder, decoder, optimizer, x_batch):
    with tf.GradientTape() as tape:
        # Encoder: produce mean and log-variance of latent distribution
        z_mean, z_log_var = encoder(x_batch, training=True)

        # Reparameterization trick: sample z ~ N(z_mean, exp(z_log_var))
        eps = tf.random.normal(tf.shape(z_mean))
        z = z_mean + tf.exp(0.5 * z_log_var) * eps

        # Decoder: reconstruct from latent sample
        x_reconstructed = decoder(z, training=True)

        # Reconstruction loss (e.g., binary cross-entropy)
        recon_loss = tf.reduce_mean(
            tf.keras.losses.binary_crossentropy(x_batch, x_reconstructed)
        )

        # KL divergence: regularize latent distribution toward N(0,1)
        kl_loss = -0.5 * tf.reduce_mean(
            1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var)
        )

        # ELBO: maximize evidence lower bound = minimize (recon + KL)
        total_loss = recon_loss + kl_loss

    # Single backward pass through both encoder and decoder
    all_vars = encoder.trainable_variables + decoder.trainable_variables
    grads = tape.gradient(total_loss, all_vars)
    optimizer.apply_gradients(zip(grads, all_vars))
    return recon_loss, kl_loss

GAN Training Step with Alternating Updates

@tf.function
def gan_train_step(generator, discriminator,
                   gen_optimizer, disc_optimizer,
                   real_batch, latent_dim=64,
                   n_disc_steps=1):
    batch_size = tf.shape(real_batch)[0]

    # ---- Phase A: Update discriminator ----
    for _ in range(n_disc_steps):
        noise = tf.random.normal([batch_size, latent_dim])
        with tf.GradientTape() as disc_tape:
            fake_samples = generator(noise, training=True)

            real_logits = discriminator(real_batch, training=True)
            fake_logits = discriminator(fake_samples, training=True)

            # Discriminator wants: real → 1, fake → 0
            disc_loss_real = tf.reduce_mean(
                tf.keras.losses.binary_crossentropy(
                    tf.ones_like(real_logits), real_logits, from_logits=True
                )
            )
            disc_loss_fake = tf.reduce_mean(
                tf.keras.losses.binary_crossentropy(
                    tf.zeros_like(fake_logits), fake_logits, from_logits=True
                )
            )
            disc_loss = disc_loss_real + disc_loss_fake

        disc_grads = disc_tape.gradient(disc_loss, discriminator.trainable_variables)
        disc_optimizer.apply_gradients(
            zip(disc_grads, discriminator.trainable_variables)
        )

    # ---- Phase B: Update generator ----
    noise = tf.random.normal([batch_size, latent_dim])
    with tf.GradientTape() as gen_tape:
        fake_samples = generator(noise, training=True)
        fake_logits = discriminator(fake_samples, training=True)

        # Generator wants discriminator to output 1 for fakes
        gen_loss = tf.reduce_mean(
            tf.keras.losses.binary_crossentropy(
                tf.ones_like(fake_logits), fake_logits, from_logits=True
            )
        )

    gen_grads = gen_tape.gradient(gen_loss, generator.trainable_variables)
    gen_optimizer.apply_gradients(zip(gen_grads, generator.trainable_variables))

    return disc_loss, gen_loss

The recurring principle across GNNs, VAEs, GANs, RL: Update scheduling is a stability control, not a style choice. Same computations in a different order can mean a different training method.

Demo: RL Loop and scikit-learn Integration

Spotlight 1: REINFORCE-Style Policy Gradient

A minimal policy gradient step — a single batch update showing the core mechanics:

import tensorflow as tf
import numpy as np

# Environment signature
OBS_DIM = 4
N_ACTIONS = 2

# Policy network
policy_net = tf.keras.Sequential([
    tf.keras.layers.Dense(8, activation="relu", input_shape=(OBS_DIM,)),
    tf.keras.layers.Dense(N_ACTIONS),  # outputs logits over actions
])
policy_optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3)

# Fake batch of observations (would come from environment in real RL)
observations = tf.random.normal([8, OBS_DIM], seed=42)
actions_taken = tf.constant([0, 1, 0, 1, 0, 1, 0, 1])  # actions that generated returns
returns = tf.constant([1.0, 0.5, 1.0, 0.5, 1.0, 0.5, 1.0, 0.5])  # discounted rewards

with tf.GradientTape() as tape:
    logits = policy_net(observations)
    log_probs = tf.nn.log_softmax(logits)  # log probabilities over actions

    # Select log probabilities of taken actions
    one_hot_actions = tf.one_hot(actions_taken, N_ACTIONS)
    selected_log_probs = tf.reduce_sum(log_probs * one_hot_actions, axis=1)

    # REINFORCE loss: maximize expected return = minimize negative
    policy_loss = -tf.reduce_mean(selected_log_probs * returns)

grads = tape.gradient(policy_loss, policy_net.trainable_variables)
policy_optimizer.apply_gradients(zip(grads, policy_net.trainable_variables))
print(f"Policy step done. Loss: {policy_loss.numpy():.4f}")
# Output: Policy step done. Loss: 0.4628

Spotlight 2: DQN-Style Update with Experience Replay

class TinyReplayBuffer:
    def __init__(self, capacity=1000):
        self.buffer = []
        self.capacity = capacity

    def add(self, state, action, reward, next_state, done):
        self.buffer.append((state, action, reward, next_state, done))
        if len(self.buffer) > self.capacity:
            self.buffer.pop(0)  # remove oldest

    def sample(self, n):
        indices = np.random.choice(len(self.buffer),
                                   size=min(n, len(self.buffer)),
                                   replace=False)
        return [self.buffer[i] for i in indices]

def epsilon_greedy(state, q_net, epsilon, n_actions):
    if state.ndim == 1:
        state = state[np.newaxis, :]
    batch_size = state.shape[0]
    roles = np.random.uniform(0, 1, size=batch_size)
    random_actions = np.random.randint(0, n_actions, size=batch_size)
    q_values = q_net(state, training=False)
    best_actions = tf.argmax(q_values, axis=1).numpy()
    return np.where(roles < epsilon, random_actions, best_actions)

# Build Q-network and target network
q_net = tf.keras.Sequential([
    tf.keras.layers.Dense(8, activation="relu", input_shape=(4,)),
    tf.keras.layers.Dense(2),  # Q-values per action
])
target_net = tf.keras.models.clone_model(q_net)
target_net.set_weights(q_net.get_weights())

q_optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3)
replay_buffer = TinyReplayBuffer(capacity=64)
EPSILON = 0.2
GAMMA = 0.99

# Fill replay buffer
for _ in range(20):
    s = np.random.normal(size=(4,)).astype(np.float32)
    a = int(epsilon_greedy(s, q_net, EPSILON, N_ACTIONS)[0])
    r = float(np.random.normal())
    s2 = (s + 0.1).astype(np.float32)
    replay_buffer.add(s, a, r, s2, False)

# One DQN-style update step
batch = replay_buffer.sample(8)
if batch:
    states = tf.constant([t[0] for t in batch])
    actions = tf.constant([t[1] for t in batch])
    rewards = tf.constant([t[2] for t in batch], dtype=tf.float32)
    next_states = tf.constant([t[3] for t in batch])

    # Target Q-values from target network
    next_q = target_net(next_states, training=False)
    max_next_q = tf.reduce_max(next_q, axis=1)
    td_targets = rewards + GAMMA * max_next_q

    with tf.GradientTape() as tape:
        all_q = q_net(states, training=True)
        # Gather Q-values for taken actions
        indices = tf.stack([tf.range(len(batch)), actions], axis=1)
        predicted_q = tf.gather_nd(all_q, indices)
        dqn_loss = tf.reduce_mean(tf.square(td_targets - predicted_q))

    grads = tape.gradient(dqn_loss, q_net.trainable_variables)
    q_optimizer.apply_gradients(zip(grads, q_net.trainable_variables))
    print(f"DQN loss: {dqn_loss.numpy():.4f}")
    # Output: DQN loss: 1.6844

Spotlight 3: Keras Model Wrapped for scikit-learn GridSearchCV

This pattern wraps a Keras classifier as a scikit-learn estimator, enabling GridSearchCV for hyperparameter tuning:

import warnings
warnings.filterwarnings("ignore")

from sklearn.model_selection import GridSearchCV
from sklearn.datasets import make_classification

def build_keras_model(units=16, lr=0.01):
    model = tf.keras.Sequential([
        tf.keras.layers.Dense(units, activation="relu", input_shape=(20,)),
        tf.keras.layers.Dense(3, activation="softmax"),
    ])
    model.compile(
        optimizer=tf.keras.optimizers.Adam(learning_rate=lr),
        loss="sparse_categorical_crossentropy",
        metrics=["accuracy"],
    )
    return model

class KerasClassifierWrapper:
    """scikit-learn compatible estimator wrapping a Keras model."""

    def __init__(self, units=16, lr=0.01, epochs=10):
        self.units = units
        self.lr = lr
        self.epochs = epochs
        self.model = None

    def get_params(self, deep=True):
        return {"units": self.units, "lr": self.lr, "epochs": self.epochs}

    def set_params(self, **params):
        for key, value in params.items():
            setattr(self, key, value)
        return self

    def fit(self, x, y):
        self.model = build_keras_model(units=self.units, lr=self.lr)
        self.model.fit(x, y, epochs=self.epochs, verbose=0)
        return self

    def predict(self, x):
        probs = self.model.predict(x)
        return np.argmax(probs, axis=1)

    def score(self, x, y):
        y_pred = self.predict(x)
        return float(np.mean(y_pred == y))

# Synthetic dataset
x_data, y_data = make_classification(n_samples=200, n_features=20, n_classes=3,
                                      n_informative=5, random_state=42)

# Grid search over units and learning rate
param_grid = {"units": [8, 16], "lr": [0.01, 0.001]}
gs = GridSearchCV(KerasClassifierWrapper(), param_grid, cv=2, scoring="accuracy")
gs.fit(x_data, y_data)

print(f"Best params: {gs.best_params_}")
print(f"Best score:  {gs.best_score_:.4f}")
# Output: Best params: {'lr': 0.01, 'units': 16}
# Output: Best score:  0.4000

What this achieves:

• Keras model participates in scikit-learn’s cross-validation pipeline.
• GridSearchCV explores hyperparameter combinations and selects the best.
• The KerasClassifierWrapper pattern generalizes: any Keras model can be wrapped and plugged into scikit-learn’s rich tooling (cross-validation, pipeline composition, hyperparameter search).

Summary: Key Principles Across All Modules

The Custom Training Mental Model

flowchart TD
    A[Dataset] --> B[GradientTape Scope]
    B --> C[Forward Pass - inside tape]
    C --> D[Loss - scalar, explicit reduction]
    D --> E[Gradients - 1 per trainable variable]
    E --> F{Controls?}
    F -->|Clipping| G[Clip before update]
    F -->|Schedule| H[Update LR per step]
    G --> I[Optimizer Update]
    H --> I
    I --> J[Logging and Validation]
    J --> K{Stable?}
    K -->|Yes| A
    K -->|No - NaN/inf| L[Fail fast]

Quick Reference: When to Use Each Tool

Production Stability Checklist

Before considering a training run production-ready:

• Loss reduces to a scalar with a justified reduction strategy (mean vs. sum)
• All trainable variables appear in model.trainable_variables
• Gradients are finite and non-None for all tracked variables
• Gradient norms are logged and clipping threshold is empirically justified
• Learning rate schedule is applied per step and verified in logs
• Mixed precision uses loss scaling when FP16 is enabled
• Custom metrics have explicit reset_state() logic at epoch boundaries
• Optimizer state is saved and restored correctly for checkpoint continuation
• Distribution strategy is chosen and configured for the target hardware
• Global batch size changes are compensated by learning rate adjustment
• At least one variable value actually changes after each apply_gradients() call

Search Terms

tensorflow · custom · optimization · deep · neural · networks · machine · data · science · loop · gradient · checklist · distributed · clipping · model · nested · components · correctness · loss · mixed · precision · execution · layer · loops